Mr. Kitayuta's Colorful Graph (Codeforces 506D)
Tags: Others, dsu

Mr. Kitayuta has just bought an undirected graph with n vertices and m edges. The vertices of the graph are numbered from 1 to n. Each edge, namely edge i, has a color ci, connecting vertex ai and bi.Mr. Kitayuta wants you to process the following q queries.In the i-th query, he gives you two integers - ui and vi.Find the number of the colors that satisfy the following condition: the edges of that color connect vertex ui and vertex vi directly or indirectly.

Input: The first line of the input contains space-separated two integers - n and m(2 ≤ n ≤ 105, 1 ≤ m ≤ 105), denoting the number of the vertices and the number of the edges, respectively.The next m lines contain space-separated three integers - ai, bi(1 ≤ ai < bi ≤ n) and ci(1 ≤ ci ≤ m). Note that there can be multiple edges between two vertices. However, there are no multiple edges of the same color between two vertices, that is, if i ≠ j, (ai, bi, ci) ≠ (aj, bj, cj).The next line contains a integer- q(1 ≤ q ≤ 105), denoting the number of the queries.Then follows q lines, containing space-separated two integers - ui and vi(1 ≤ ui, vi ≤ n). It is guaranteed that ui ≠ vi.

Output: For each query, print the answer in a separate line.

Note: Let's consider the first sample.     The figure above shows the first sample.    Vertex 1 and vertex 2 are connected by color 1 and 2.  Vertex 3 and vertex 4 are connected by color 3.  Vertex 1 and vertex 4 are not connected by any single color. 

Sample Inputs:
4 5
1 2 1
1 2 2
2 3 1
2 3 3
2 4 3
3
1 2
3 4
1 4

5 7
1 5 1
2 5 1
3 5 1
4 5 1
1 2 2
2 3 2
3 4 2
5
1 5
5 1
2 5
1 5
1 4

Sample Outputs:
2
1
0

1
1
1
1
2

Solution (Python):
import os, sys, collections, heapq, itertools
if os.path.exists('in.txt'):
    file = open('in.txt')
    input = file.readline
else:
    input = sys.stdin.readline
from functools import reduce
from itertools import accumulate, chain, combinations, count
from itertools import groupby, permutations, product, zip_longest
from heapq import heapify, heappush, heappop, heapreplace, merge, nlargest, nsmallest
from collections import Counter, defaultdict, deque
from bisect import bisect_left, bisect_right
from math import comb, floor, ceil, inf, gcd, sqrt
import functools
readi = lambda : int(input())
readis = lambda : list(map(int, input().split()))
reads = lambda : input().rstrip()
readss = lambda : input().split()

class UF:

    def __init__(self, n):
        self.n = n
        self.FA = list(range(n))

    def getroot(self, x):
        path = []
        while x != self.FA[x]:
            path.append(x)
            x = self.FA[x]
        for w in path:
            self.FA[w] = x
        return x

    def union(self, x, y):
        self.FA[self.getroot(x)] = self.getroot(y)

    def getgroup(self, vertices):
        G = defaultdict(list)
        for i in vertices:
            G[self.getroot(i)].append(i)
        return G

def solution():
    (n, m) = readis()
    E = defaultdict(list)
    for _ in range(m):
        (x, y, c) = readis()
        E[c].append((x, y))
    u = UF(n + 1)
    ST = [0] * (n + 1)
    P = defaultdict(int)
    (L, cnt) = (int(sqrt(m)) // 10, 0)
    for es in E.values():
        for (x, y) in es:
            u.union(x, y)
        st = set(chain(*es))
        G = u.getgroup(st)
        for v in G.values():
            if len(v) <= L:
                for (x, y) in combinations(sorted(v), 2):
                    P[y * n + x] += 1
            else:
                for x in v:
                    ST[x] |= 1 << cnt
                cnt += 1
        for x in st:
            u.FA[x] = x
    q = readi()
    for _ in range(q):
        (x, y) = readis()
        if x > y:
            (x, y) = (y, x)
        print(P[y * n + x] + bin(ST[x] & ST[y]).count('1'))
for i in range(1):
    solution()

Solution (Python):
import os, sys, collections, heapq, itertools
if os.path.exists('in.txt'):
    file = open('in.txt')
    input = file.readline
else:
    input = sys.stdin.readline
from functools import reduce
from itertools import accumulate, chain, combinations, count
from itertools import groupby, permutations, product, zip_longest
from heapq import heapify, heappush, heappop, heapreplace, merge, nlargest, nsmallest
from collections import Counter, defaultdict, deque
from bisect import bisect_left, bisect_right
from math import comb, floor, ceil, inf, gcd, sqrt
import functools
readi = lambda : int(input())
readis = lambda : list(map(int, input().split()))
reads = lambda : input().rstrip()
readss = lambda : input().split()

class UF:

    def __init__(self, n):
        self.n = n
        self.FA = list(range(n))

    def getroot(self, x):
        (root, FA) = (x, self.FA)
        while root != FA[root]:
            root = FA[root]
        while x != root:
            (FA[x], x) = (root, FA[x])
        return root

    def union(self, x, y):
        self.FA[self.getroot(x)] = self.getroot(y)

    def getgroup(self, vertices):
        G = defaultdict(list)
        for i in vertices:
            G[self.getroot(i)].append(i)
        return G

def solution():
    (n, m) = readis()
    E = defaultdict(list)
    for _ in range(m):
        (x, y, c) = readis()
        E[c].append((x, y))
    u = UF(n + 1)
    ST = [0] * (n + 1)
    P = defaultdict(int)
    (L, cnt) = (int(sqrt(m)) // 10, 0)
    for es in E.values():
        for (x, y) in es:
            u.union(x, y)
        st = set(chain(*es))
        G = u.getgroup(st)
        for v in G.values():
            if len(v) <= L:
                for (x, y) in combinations(sorted(v), 2):
                    P[y * n + x] += 1
            else:
                for x in v:
                    ST[x] |= 1 << cnt
                cnt += 1
        for x in st:
            u.FA[x] = x
    q = readi()
    for _ in range(q):
        (x, y) = readis()
        if x > y:
            (x, y) = (y, x)
        print(P[y * n + x] + bin(ST[x] & ST[y]).count('1'))
for i in range(1):
    solution()

Solution (Python):
import os, sys, collections, heapq, itertools
if os.path.exists('in.txt'):
    file = open('in.txt')
    input = file.readline
else:
    input = sys.stdin.readline
from functools import reduce
from itertools import accumulate, chain, combinations, count
from itertools import groupby, permutations, product, zip_longest
from heapq import heapify, heappush, heappop, heapreplace, merge, nlargest, nsmallest
from collections import Counter, defaultdict, deque
from bisect import bisect_left, bisect_right
from math import comb, floor, ceil, inf, gcd, sqrt
import functools
readi = lambda : int(input())
readis = lambda : list(map(int, input().split()))
reads = lambda : input().rstrip()
readss = lambda : input().split()

class UF:

    def __init__(self, n):
        self.n = n
        self.FA = list(range(n))

    def getroot(self, x):
        path = []
        FA = self.FA
        while x != FA[x]:
            path.append(x)
            x = FA[x]
        for w in path:
            FA[w] = x
        return x

    def union(self, x, y):
        self.FA[self.getroot(x)] = self.getroot(y)

    def getgroup(self, vertices):
        G = defaultdict(list)
        for i in vertices:
            G[self.getroot(i)].append(i)
        return G

def solution():
    (n, m) = readis()
    E = defaultdict(list)
    for _ in range(m):
        (x, y, c) = readis()
        E[c].append((x, y))
    u = UF(n + 1)
    ST = [0] * (n + 1)
    P = defaultdict(int)
    (L, cnt) = (int(sqrt(m)) // 10, 0)
    for es in E.values():
        for (x, y) in es:
            u.union(x, y)
        st = set(chain(*es))
        G = u.getgroup(st)
        for v in G.values():
            if len(v) <= L:
                for (x, y) in combinations(sorted(v), 2):
                    P[y * n + x] += 1
            else:
                for x in v:
                    ST[x] |= 1 << cnt
                cnt += 1
        for x in st:
            u.FA[x] = x
    q = readi()
    for _ in range(q):
        (x, y) = readis()
        if x > y:
            (x, y) = (y, x)
        print(P[y * n + x] + bin(ST[x] & ST[y]).count('1'))
for i in range(1):
    solution()