【Python3】 AtCoder Beginner Contest 120 D – Decayed Bridges

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)


この記事を書いた人

サトゥー

東大学際情報学府M1。情報科学と教養の海に溺れています。面白いことをやるのがすきです。