Codeforces Round 891 (Div. 3) G题解

Given a tree consisting of nn vertices. A tree is a connected undirected graph without cycles. Each edge of the tree has its weight, wiwi.

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 SS.
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 998244353998244353.

代码：

int n, m, fa[N], sz[N];
struct Edge {
    int x, y, v;
    Edge(int x = 0, int y = 0, int v = 0) : x(x), y(y), v(v) {}
    bool operator < (const Edge &A) const {
        return v == A.v ? (x == A.x ? y < A.y : x < A.x) : v < A.v;
    }
} a[N]; // 建图
 
int Mul(int x, int y) {
	return 1ll * x * y % MOD2;
} //乘法运算
 
int qpow(int x, long long y) {
	// cerr << x << " " << y << endl;
	int res = 1;
	for(; y; x = Mul(x, x), y >>= 1)
		if(y & 1)
			res = Mul(res, x);
	return res % MOD2;
} //快速幂
 
int find(int x) {
    return x == fa[x] ? fa[x] : fa[x] = find(fa[x]);
} // 找祖宗
 
int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(NULL), std::cout.tie(NULL);       
 
    int t;
    std::cin >> t;
 
    while(t--) {
    std::cin >> n >> m;
    for(int i = 1; i < n; i++)   
        std::cin >> a[i].x >> a[i].y >> a[i].v; // 把起点、终点、边权存入
    
    std::sort(a + 1, a + n); // 根据边权大小进行排序
 
    for(int i = 1; i <= n; i++)
        fa[i] = i, sz[i] = 1; // 并查集初始化
 
    LL ans = 1;
    for(int i = 1; i < n; i++) {
        int x = find(a[i].x), y = find(a[i].y), v = a[i].v;
        LL cnt = 1ll * sz[x] * sz[y] - 1; //所有有可能加的边的数量大小
        fa[x] = y;
        sz[y] += sz[x];
        ans = Mul(ans, qpow(m - v + 1, cnt)); //乘法原理
    }
 
    std::cout << ans << "\n";
    }
}