ABC 049 D 連結

競技プログラミング

qiita.com

続きです。

atcoder.jp

考察まとめです。

Union-Find木を使って解いてみる。まずTLEになったやつ。

class UnionFind():
    def __init__(self, n=None):
        if type(n) == int:
            self.par = [i for i in range(n)]
            self.rank = [0 for i in range(n)]
        else:
            self.par = []
            self.rank = []

    def init(self, n):
        self.par = [i for i in range(n)]
        self.rank = [0 for i in range(n)]

    def find(self, x):
        """ ノード x の根を求める """
        if self.par[x] == x:
            return x
        self.par[x] = self.find(self.par[x])
        return self.par[x]

    def unite(self, x, y):
        root_x = self.find(x); root_y = self.find(y)
        if root_x == root_y:
            return
        if self.rank[root_x] < self.rank[root_y]:
            self.par[root_x] = root_y
        elif self.rank[root_y] < self.rank[root_x]:
            self.par[root_y] = root_x
        else:
            self.par[root_y] = root_x
            self.rank[root_x] += 1

    def same(self, x, y):
        return self.find(x) == self.find(y)

N, K, L = map(int, input().split())
road = UnionFind(N)
rail = UnionFind(N)
for i in range(K):
    p, q = map(int, input().split())
    road.unite(p-1, q-1)
for j in range(L):
    r, s = map(int, input().split())
    rail.unite(r-1, s-1)

for i in range(N):
    for j in range(N):
        if road.find(i) == road.find(j) and rail.find(i) == rail.find(j):
            ans[i] += 1
for i in range(N):
    print(ans[i], end=' ')
print('')

これは、結果的に

for i in range(N):
    for j in range(N):
        if road.find(i) == road.find(j) and rail.find(i) == rail.find(j):
            ans[i] += 1

この部分がまずかった。

これは、数え上げで、計算量を減らすことができるようだ。

for i in range(N):
    ans[(road.find(i), rail.find(i))] += 1

for i in range(N):
    print(ans[(road.find(i), rail.find(i))], end=' ')
print('')

こんな感じに。

これは、各頂点が、少なくとも 1 つのグループに属しているので、根が同じものを数えれば、どれだけの頂点数があるのか数えることができるということを使用している。

つまり、

# 考察のため、キーを一次元にした
for i in range(N):
    for j in range(N):
        if road.find(i) == road.find(j):
            ans[i] += 1

これは、

for i in range(N):
    ans[road.find(i)] += 1

とすることができる。

これを、2 次元ぽくすると、上記の計算量を落とした解答にありつける。