Union-Find を使う問題でまだ発想に時間がかかったので、ここにまとめておきます。
$N$ 個の島と $M$ 本の橋があります。$i$ 番目の橋は $A_i$ 番目の島と $B_i$ 番目の島を繋いでおり、双方向に行き来可能です。
はじめ、どの2つの島についてもいくつかの橋を渡って互いに行き来できます。
調査の結果、老朽化のためこれら $M$ 本の橋は1番目の橋から順に全て崩落することがわかりました。
「いくつかの橋を渡って互いに行き来できなくなった2つの島の組 $(a, b) (a < b)$ の数」を不便さと呼ぶことにします。
各 $i (1 \leq i \leq M)$ について、$i$ 番目の橋が崩落した直後の不便さを求めてください。
https://atcoder.jp/contests/abc120/tasks/abc120_d
島 $a$ と島 $b$ がつながっているかどうかを判定すればよさそうなので、 Union-Find を使うのが有効そう。しかし、 Union-Find は要素をまとめることはできても、今回のように「橋が崩落して要素が切れてしまう」ことは表現できません。
そこで、時間経過を逆転させて考えてみます。つまり、橋が崩壊していくのではなく、 $i = M, M-1, ……, 1$ の順で橋が復旧していくときの答えの変化を順に考えます。こうすることで Union-Find をうまく使える形に帰着させることができました。
あるデータ構造が使えないように見える問題形式でも、問題の見方を変えることでこのように活用することができます。
では、逆順に更新するということで、はじめに ans[M-1]
(最終的な状態)を考えますと、次のようになります。
ans[M-1] = N*(N-1)//2
そして、index を減らしながら更新していきます。橋 $i+1$ が復旧すると、group_size(A_{i+1}) * group_size(B_{i+1})
分だけ答えが減ることに注意して、次のように更新できます。
if (A_{i+1} と B_{i+1} がつながっている):
ans[i] = ans[i+1]
else:
ans[i] = ans[i+1] - group_size(A_{i+1}) * group_size(B_{i+1})
結局、回答は次のようになります。
class UnionFind:
def __init__(self, n):
self.n = n
self.parent = [i for i in range(n)] # 親
self.rank = [1] * n # 木の高さ
self.size = [1] * n # size[i] は i を根とするグループのサイズ
def find(self, x): # x の根を返す
if self.parent[x] == x:
return x
else:
self.parent[x] = self.find(self.parent[x]) # 経路圧縮
return self.parent[x]
def unite(self, x, y): # x, y の属する集合を併合する
x = self.find(x)
y = self.find(y)
if x != y:
if self.rank[x] < self.rank[y]:
self.parent[x] = y
self.size[y] += self.size[x]
else:
self.parent[y] = x
self.size[x] += self.size[y]
if self.rank[x] == self.rank[y]:
self.rank[x] += 1
def is_same(self, x, y): # x, y が同じ集合に属するか判定する
return self.find(x) == self.find(y)
def group_size(self, x): # x が属する集合の大きさを返す
return self.size[self.find(x)]
def __str__(self): # print 表示用
return '\n'.join('{}: {}'.format(r, self.group_members(r)) for r in self.roots())
N, M = map(int, input().split())
bridge = [list(map(int, input().split())) for _ in range(M)]
ans = [0]*M
uf = UnionFind(N)
# うしろからつないでいく
ans[M-1] = N*(N-1)//2
for i in range(M-2, -1, -1):
A, B = bridge[i+1]
if uf.is_same(A-1, B-1):
ans[i] = ans[i+1]
else:
ans[i] = ans[i+1] - uf.group_size(A-1) * uf.group_size(B-1)
uf.unite(A-1, B-1)
for a in ans:
print(a)