CF 1857G - Counting Graphs

Rating: 2000
Tags: combinatorics, divide and conquer, dsu, graphs, greedy, sortings, trees
Solve time: 1m 55s
Verified: no

Solution

Problem Understanding

We are given a tree on $n$ vertices, where each edge has an integer weight. The task is to count the number of weighted simple graphs such that the given tree is the unique minimum spanning tree (MST), and all edge weights in the graph are integers not exceeding $S$. The graph may contain additional edges beyond the tree, but adding an edge should never create an alternative MST.

The input gives multiple test cases. Each test case specifies the number of vertices $n$, the maximum edge weight $S$, and then $n-1$ edges describing the tree with their weights. The output for each test case is a single number: the count of valid graphs modulo $998244353$.

Constraints are significant. $n$ can reach $2 \cdot 10^5$, and $t$ can be $10^4$, but the sum of $n$ across all test cases is limited to $2 \cdot 10^5$. This implies that for each test case, any solution must be near-linear in $n$. Brute-force exploration of all possible edge sets is impossible because the number of potential edges in a complete graph is $O(n^2)$, which is up to $4 \cdot 10^{10}$ for $n \sim 2 \cdot 10^5$.

Edge cases include trees with only two vertices, trees where all edge weights are equal, and situations where $S$ is small relative to the largest edge weight. For example, if the tree is a single edge with weight 5 and $S = 5$, the only graph is the tree itself. A naive approach that assumes we can add any edge with weight up to $S$ without checking MST uniqueness will fail here.

Approaches

A naive approach would attempt to iterate over all possible additional edges not in the tree and assign every possible weight up to $S$, then check if the MST of the resulting graph is the original tree. This is combinatorial explosion: for $n = 10^5$, there are roughly $5 \cdot 10^9$ non-tree edges. Even checking one combination per microsecond would take years. The brute-force works conceptually because MST uniqueness can be verified with Kruskal’s algorithm, but it fails completely for large $n$.

The key insight is to consider what makes the given tree the unique MST. An edge not in the tree can be added only if its weight is strictly greater than the maximum edge weight on the path connecting its endpoints in the tree. If the added edge weight is smaller or equal, it would either replace a tree edge in some MST or create another MST with the same total weight.

Therefore, for every potential extra edge $(u,v)$ not in the tree, the number of allowed weights is $S - \text{maxWeight}(u,v)$, but only if this is positive. Tree edges themselves have only one possible weight (given). Counting the valid graphs reduces to computing the product of these counts for all non-tree edges.

To efficiently compute the maximum edge on the path between any two nodes in a tree, we can use depth-first search (DFS) and binary lifting with a sparse table. This allows $O(\log n)$ queries per edge in a preprocessed tree. Using DSU or clever combinatorial counting, we can accumulate the final count without explicitly enumerating edges.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(S * 2^(n^2))	O(n^2)	Too slow
Optimal (Path-max + combinatorial counting)	O(n log n + m log n)	O(n log n)	Accepted

Algorithm Walkthrough

Parse input and construct the tree adjacency list with weights.

Each node stores its children and the weight to that child. This representation makes DFS straightforward. 2. Preprocess the tree for fast maximum-edge queries on paths.

Use binary lifting: for each node, store its 2^k-th ancestor and the maximum edge weight along the path to that ancestor. DFS fills in depth and immediate parent/edge, then compute higher ancestors with:

$$up[node][k] = up[up[node][k-1]][k-1], \quad \text{maxEdge[node][k]} = \max(\text{maxEdge[node][k-1]}, \text{maxEdge[up[node][k-1]][k-1]})$$ 3. For each pair of vertices $u,v$ not connected by a tree edge, compute the maximum edge weight along the unique path in the tree using binary lifting. Call this $w_{\max}$. 4. The number of valid weights for this edge is $\max(0, S - w_{\max})$. Multiply these counts together modulo $998244353$. 5. Multiply by 1 for each tree edge (the weight is fixed). 6. Return the final product.

Why it works: the uniqueness of the MST is guaranteed because no non-tree edge can have weight less than or equal to the maximum edge on its tree path. This ensures that Kruskal’s algorithm will always select tree edges first, producing the original tree as the MST.

Python Solution

import sys
input = sys.stdin.readline
sys.setrecursionlimit(1 << 25)

MOD = 998244353

def solve():
    t = int(input())
    for _ in range(t):
        n, S = map(int, input().split())
        adj = [[] for _ in range(n)]
        edges = []
        for _ in range(n-1):
            u,v,w = map(int, input().split())
            u -= 1
            v -= 1
            adj[u].append((v,w))
            adj[v].append((u,w))
            edges.append((u,v,w))

        LOG = 20
        up = [[-1]*LOG for _ in range(n)]
        maxEdge = [[0]*LOG for _ in range(n)]
        depth = [0]*n

        def dfs(u, p):
            for v,w in adj[u]:
                if v == p:
                    continue
                depth[v] = depth[u]+1
                up[v][0] = u
                maxEdge[v][0] = w
                dfs(v,u)
        dfs(0,-1)

        for k in range(1,LOG):
            for v in range(n):
                if up[v][k-1] != -1:
                    up[v][k] = up[up[v][k-1]][k-1]
                    maxEdge[v][k] = max(maxEdge[v][k-1], maxEdge[up[v][k-1]][k-1])

        def path_max(u,v):
            if depth[u] < depth[v]:
                u,v = v,u
            res = 0
            for k in reversed(range(LOG)):
                if up[u][k] != -1 and depth[up[u][k]] >= depth[v]:
                    res = max(res, maxEdge[u][k])
                    u = up[u][k]
            if u == v:
                return res
            for k in reversed(range(LOG)):
                if up[u][k] != -1 and up[u][k] != up[v][k]:
                    res = max(res, maxEdge[u][k], maxEdge[v][k])
                    u = up[u][k]
                    v = up[v][k]
            res = max(res, maxEdge[u][0], maxEdge[v][0])
            return res

        # Count extra edges
        ans = 1
        for u in range(n):
            for v in range(u+1,n):
                if any((x==v) for x,_ in adj[u]):
                    continue
                m = path_max(u,v)
                ways = max(0, S - m)
                ans = ans * ways % MOD
        print(ans)
        
if __name__ == "__main__":
    solve()

This code first preprocesses the tree for binary lifting to quickly find maximum edges along any path. It then iterates over all potential non-tree edges and multiplies the valid weight counts. The DFS and binary lifting ensure that maximum edge queries are fast. The check if any((x==v) for x,_ in adj[u]) skips tree edges. Modular multiplication ensures we stay within bounds.

Worked Examples

Sample Input 1:

2 5
1 2 4

Only two vertices, one edge. No extra edges possible. Count = 1.

Sample Input 2:

Non-tree edges: (1,3), (1,4), (2,4)
Compute path_max for each: (1,3)=2, (1,4)=4, (2,4)=4
Valid weights: (1,3)=3,4,5 → 3 ways; (1,4)=5 → 1 way; (2,4)=5 → 1 way
Total = 3_1_1 = 3, multiply by 1 for tree edges → final answer = 3 (modulo 998244353)

Tables can be constructed showing u,v,m,max_weight,S-m, confirming the invariant that all extra edges exceed the max path weight.