ABC040D

最近キックボクシングジムに通い始めました。

ko_ya346です。

Atcoder過去問埋めをやっていきます。

問題はこちら。

atcoder.jp

何を考えた？

道で繋がっている都市の個数を数えるのでUnionFind木を使えば楽そうです。

class UnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def size(self, x):
        return -self.parents[self.find(x)]

次は入力

道とクエリは新しい順にソートします。

クエリ順に出力する必要があるので、クエリのindexを記録しています（[i]のとこ）。

N, M = map(int, input().split())
road = sorted([list(map(int, input().split())) for _ in range(M)], key=lambda x:x[2], reverse=True)
Q = int(input())

uf = UnionFind(N)
ans = [None]*Q
query = sorted([list(map(int, input().split())) + [i] for i in range(Q)], key=lambda x:x[1], reverse=True)

UnionFind木は簡単に辺を増やせますが辺を削除することは難しいらしいです。

なのでqueryの住民が要求する橋を新しい順にソートし、その住民を満足させる辺を足すことにします。

同じ道を何度も足すのは無駄なので、前回チェックしたindexをr_indに保持し、クエリが始まる度に続きからチェックするようにします。

毎回道を全てチェックするとTLEになるのでご注意ください（3敗）

住民が満足する道を全て足したところで、その住民がいる地点のグループ数を計算します。

#前回追加したところまでのindexを保持して
#木の追加に重複がないようにする
r_ind = 0
for v, w, i in query:
    while r_ind < M and road[r_ind][2] > w:
        a, b, y = road[r_ind]
        uf.union(a-1, b-1)
        r_ind += 1
    ans[i] = uf.size(v-1)

print(*ans, sep="\n")

全体の解法はこちら

class UnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def size(self, x):
        return -self.parents[self.find(x)]

N, M = map(int, input().split())
road = sorted([list(map(int, input().split())) for _ in range(M)], key=lambda x:x[2], reverse=True)
Q = int(input())

uf = UnionFind(N)
ans = [None]*Q
query = sorted([list(map(int, input().split())) + [i] for i in range(Q)], key=lambda x:x[1], reverse=True)
# print(query)

#前回追加したところまでのindexを保持して
#木の追加に重複がないようにする
r_ind = 0
for v, w, i in query:
    while r_ind < M and road[r_ind][2] > w:
        a, b, y = road[r_ind]
        uf.union(a-1, b-1)
        r_ind += 1
    ans[i] = uf.size(v-1)

print(*ans, sep="\n")

UnionFind木の構造は複雑ですが、使いこなせるとすごく便利ですね！