Counting Graphs (Codeforces 1857G)
Tags: Greedy, divide and conquer

Given a tree consisting of $$$n$$$ vertices. A tree is a connected undirected graph without cycles. Each edge of the tree has its weight, $$$w_i$$$.Your task is to count the number of different graphs that satisfy all four conditions:  The graph does not have self-loops and multiple edges.  The weights on the edges of the graph are integers and do not exceed $$$S$$$.  The graph has exactly one minimum spanning tree.  The minimum spanning tree of the graph is the given tree. Two graphs are considered different if their sets of edges are different, taking into account the weights of the edges.The answer can be large, output it modulo $$$998244353$$$.

Input: The first line contains an integer $$$t$$$ ($$$1\le t\le 10^4$$$) — the number of test cases.The first line of each test case contains two integers $$$n$$$ and $$$S$$$ ($$$2 \le n \le 2 \cdot 10^5, 1\le S\le 10^9$$$) — the number of vertices and the upper bound of the weights.Then follow $$$n-1$$$ lines describing the tree, the $$$i$$$-th line contains three integers $$$u_i$$$, $$$v_i$$$, and $$$w_i$$$ ($$$1\le u_i,v_i\le n, u_i \ne v_i, 1\le w_i\le S$$$) — an edge in the tree with weight $$$w_i$$$.It is guaranteed that the sum of $$$n$$$ for all tests does not exceed $$$2\cdot 10^5$$$.

Output: For each test, output the number of different graphs that satisfy the conditions, modulo $$$998244353$$$.

Note: In the first sample, there is only one graph, which is the given tree.In the second samle, the given tree looks like this:    All possible graphs for the second sample are shown below, the minimum spanning tree is highlighted in red:  

Sample Inputs:
4
2 5
1 2 4
4 5
1 2 2
2 3 4
3 4 3
5 6
1 2 3
1 3 2
3 4 6
3 5 1
10 200
1 2 3
2 3 33
3 4 200
1 5 132
5 6 1
5 7 29
7 8 187
7 9 20
7 10 4

Sample Outputs:
1
8
80
650867886

Solution (Python):
import sys
input = sys.stdin.readline
print = sys.stdout.write
mod = 998244353

def find_set(v):
    stack = [v]
    while dsu[v] != v:
        stack.append(dsu[v])
        v = stack[-1]
    while stack:
        dsu[stack[-1]] = dsu[v]
        v = stack.pop()
    return dsu[v]
for _ in range(int(input())):
    (n, mxw) = map(int, input().split())
    l = sorted([tuple(map(int, input().split())) for _ in range(n - 1)], key=lambda x: x[2])
    ans = 1
    dsu = [i for i in range(n)]
    cnt = [1] * n
    for (v1, v2, w) in l:
        v1 -= 1
        v2 -= 1
        if find_set(v1) != find_set(v2):
            ans = ans * pow(mxw - w + 1, cnt[dsu[v1]] * cnt[dsu[v2]] - 1, mod) % mod
            cnt[dsu[v1]] += cnt[dsu[v2]]
            cnt[dsu[v2]] = 0
            dsu[v2] = dsu[v1]
    print(str(ans) + '\n')

Solution (Python):
import sys
import math
import collections
from heapq import heappush, heappop
input = sys.stdin.readline
ints = lambda : list(map(int, input().split()))

class UnionFind:

    def __init__(self, n):
        self.link = [i for i in range(n)]
        self.size = [1 for i in range(n)]

    def find(self, a):
        while self.link[a] != a:
            a = self.link[a]
        return a

    def same(self, a, b):
        return self.find(a) == self.find(b)

    def join(self, a, b):
        a = self.find(a)
        b = self.find(b)
        if self.size[a] > self.size[b]:
            (a, b) = (b, a)
        self.link[a] = b
        self.size[b] += self.size[a]
        return b
t = int(input())
for _ in range(t):
    (n, s) = ints()
    edges = []
    for _ in range(n - 1):
        (u, v, w) = ints()
        u -= 1
        v -= 1
        edges.append((u, v, w))
    edges.sort(key=lambda x: x[2])
    ans = 1
    dsu = UnionFind(n)
    for (u, v, w) in edges:
        sz1 = dsu.size[dsu.find(u)]
        sz2 = dsu.size[dsu.find(v)]
        tot = sz1 * sz2 - 1
        base = s - w + 1
        ans *= pow(base, tot, 998244353)
        ans %= 998244353
        dsu.join(u, v)
    print(ans)

Solution (Python):
import sys
import math
import collections
from heapq import heappush, heappop
input = sys.stdin.readline
ints = lambda : list(map(int, input().split()))

class UnionFind:

    def __init__(self, n):
        self.link = [i for i in range(n)]
        self.size = [1 for i in range(n)]

    def find(self, a):
        while self.link[a] != a:
            a = self.link[a]
        return a

    def same(self, a, b):
        return self.find(a) == self.find(b)

    def join(self, a, b):
        a = self.find(a)
        b = self.find(b)
        if self.size[a] > self.size[b]:
            (a, b) = (b, a)
        self.link[a] = b
        self.size[b] += self.size[a]
        return b
t = int(input())
for _ in range(t):
    (n, s) = ints()
    edges = []
    for _ in range(n - 1):
        (u, v, w) = ints()
        u -= 1
        v -= 1
        edges.append((u, v, w))
    edges.sort(key=lambda x: x[2])
    ans = 1
    dsu = UnionFind(n)
    for (u, v, w) in edges:
        sz1 = dsu.size[dsu.find(u)]
        sz2 = dsu.size[dsu.find(v)]
        tot = sz1 * sz2 - 1
        base = s - w + 1
        ans *= pow(base, tot, 998244353)
        ans %= 998244353
        dsu.join(u, v)
    print(ans)