Counting Graphs
Contest: Codeforces Round 891 (Div. 3)
Difficulty rating: 2000
Tags: combinatorics, divide and conquer, dsu, graphs, greedy, sortings, trees

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:

1. The graph does not have self-loops and multiple edges.
2. The weights on the edges of the graph are integers and do not exceed $$$S$$$.
3. The graph has exactly one minimum spanning tree.
4. 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 format:
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 format:
For each test, output the number of different graphs that satisfy the conditions, modulo $$$998244353$$$.

Constraints:
Time limit: 2.0 s
Memory limit: 256.0 MB

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:

Example input:
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

Example output:
1
8
80
650867886

Editorial:
The first observation is that the graphs will consist of $$$n$$$ vertices because the MST is fixed. Hence, the graphs will look like the given tree with some new vertices connected.
The next step is to determine the possible weights of a new edge between vertices $$$u$$$ and $$$v$$$. Let $$$P(u,v)$$$ be the maximum weight on the simple path from $$$u$$$ to $$$v$$$. I assume that we can add a new edge between vertices $$$u$$$ and $$$v$$$ with any weight in the range $$$[P(u,v)+1,S]$$$. This becomes clear when we consider some examples. If the edge has a weight less or equal than $$$P(u,v)$$$, the MST will change by taking the new edge instead of the edge with the maximal weight on the path.
Notice that if we add a new edge, and it doesn't affect the MST, we can add one more edge independently from the previous.
So now the task is to calculate $$$\prod\limits_{1\le u < v \le n}{} (S-P(u,v)+1)$$$, because for each pair of vertices, we can assign a new weight from the range, which gives us $$$S-P(u,v)$$$ cases, or we can choose not to add any edge, which is one more case.
Now, let's discuss how to calculate the formula efficiently:
1. Sort the given edges in ascending order according to their weights: $$$w_1\le w_2\le \dots \le w_{n-1}$$$.
2. We'll begin from the graph without edges, and add new ones step by step. Suppose, we already added all the edges up to $$$u_i,v_i,w_i$$$. Now, we want to add the $$$i$$$-th one. Notice that $$$w_i$$$ is greater than any of the weights before, and $$$u_i$$$ and $$$v_i$$$ are from different components.
3. After adding the edge, we need to calculate the number of paths that go through this edge. If we know the sizes of the components containing $$$u_i$$$ and $$$v_i$$$, denoted as $$$s_u$$$ and $$$s_v$$$ respectively, then there exist $$$s_u\cdot s_v-1$$$ paths through edge $$$u_i,v_i$$$ without including the path formed by these two vertices.
4. We know the number of paths, and on each of these paths, we can determine the weight we can put on the edge. Thus, we need to multiply answer by $$$(S-w_i+1)^{s_u\cdot s_v-1}$$$, using binary exponentiation.
To add edges and find the sizes of the components efficiently, you can use DSU (Disjoint Set Union).
The complexity of this approach is $$$O(n\log n)$$$.