F - MST Query Editorial by en_translator
Let \(W=10\) be the maximum edge weight. Let \(G_k (0\leq k\leq W)\) be the (unweighted) subgraph of \(G\) consisting of the edges of \(G\) with weights not more than \(k\). Let \(H\) be the number of connected components of a graph \(H\), then the sum of weights of edges in an MST (Minimum Spanning Tree) of graph \(G\) equals \(\displaystyle \sum_{k=0}^{W-1}\left( c(G_k)-1\right)\).
(Proof)
(Approach 1)
Among the edges in an MST of \(G\), there are \(N-c(G_k)\) edges whose weight is not greater than \(k\) (because it is optimal to pick lighter edges as much as possible to the MST, by Kruskal’s algorithm). Therefore, there are \(\left(N-c(G_k)\right)-\left(N-c(G_{k-1})\right)=c(G_{k-1})-c(G_k)\) edges whose weight is exactly \(k\), so the total edge weight in an MST of graph \(G\) is \(\displaystyle \sum_{k=1}^{W}k\left(c(G_{k-1})-c(G_k)\right)=\displaystyle \sum_{k=0}^{W-1}\left( c(G_k)\right)-W\cdot c(G_W)=\displaystyle \sum_{k=0}^{W-1}\left( c(G_k)-1\right)\).
(Approach 2) Among the edges in an MST of \(G\), there are \(c(G_{k-1})-1\) edges whose weight is not greater than \(c(G_{k-1})-1\) (because we need to add \(c(G_{k-1})-1\) edges to make \(G_{k-1}\) connected). Thus, by the “preposterous trick” the total edge weight in an MST of graph \(G\) is \(\displaystyle\sum_{k=1}^{W}\left(c(G_{k-1})-1\right)=\displaystyle \sum_{k=0}^{W-1}\left(c(G_k)-1\right)\). (End of proof)
For each \(k(0\leq k\lt W)\), one can manage the number of connected components of \(G_k\) with a disjoint set union (DSU) to solve this problem in a total of \(O(W(N+Q)\alpha(N))\). (The complexity can be reduced to \(O((WN+Q)\alpha(N))\).)
In the sample code below, we apply the following tricks:
- We assume that \(G\) originally has \((N-1)\) edges of weight \(10\), edges \(\{1,2\},\{1,3\},\dots,\{1,N\}\), and there are \(N-1+Q\) queries.
- Maintain the answer itself. Initially, \(c(G_{k})=N(0\leq k\lt 10)\), so \(\mathrm{ans}=10N-10\). For each query, update the answer whenever the number of connected component decreases.
from atcoder.dsu import DSU
N, Q = map(int, input().split())
uf = [DSU(N + 1) for i in range(10)]
ans = 10 * N - 10
for i in range(N - 1 + Q):
a, b, c = map(int, input().split())
for j in range(c, 10):
if not uf[j].same(a, b):
ans -= 1
uf[j].merge(a, b)
if i >= N - 1:
print(ans)
posted:
last update: