クルトンのプログラミング教室

Pythonの使い方やPythonを使った競技プログラミングの解法などを解説しています。

Pythonで理解する蟻本「2-4 食物連鎖（POJ 1182）」(p.85)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-4 食物連鎖（POJ 1182）」(p.85)
のコードをPythonで書き直したものとなっています。

入力

N\,K\\T_1\,…\,T_K\\X_1\,…\,X_K\\Y_1\,…\,Y_K

入力例



100 7
1 2 2 2 1 2 1
101 1 2 3 1 3 5
1 2 3 3 3 1 5


解答

# 入力（Tは情報のタイプ）
N, K = map(int,input().split())
T = list(map(int,input().split()))
X = list(map(int,input().split()))
Y = list(map(int,input().split()))

# Union-Find木のソースコード
class UnionFindTree():
    # n要素で初期化
    def __init__(self, n):
        self.n = n
        self.par = list(range(n))
        self.rank = [0] * n

    # 木の根を求める
    def find(self, x):
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    # xとyの属する集合を併合
    def unite(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return

        if self.rank[x] < self.rank[y]:
            self.par[x] = y
        else:
            self.par[y] = x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

    # xとyが同じ集合に属するか否か
    def same(self, x, y):
        return self.find(x) == self.find(y)

# Union-Find木を初期化
# x, x + N, x + 2 * N を x-A, x-B, x-Cの要素とする
UFT = UnionFindTree(N * 3)

ans = 0
for i in range(K):
    t = T[i]
    x = X[i] - 1    # 0, ..., N - 1の範囲にする
    y = Y[i] - 1
    
    # 正しくない番号
    if x < 0 or N <= x or y < 0 or N <= y:
        ans += 1
        continue
    
    if t == 1:
        # 「xとyは同じ種類」という情報
        if UFT.same(x, y + N) or UFT.same(x, y + 2 * N):
            ans += 1
        else:
            UFT.unite(x, y)
            UFT.unite(x + N, y + N)
            UFT.unite(x + N * 2, y + N * 2)
    else:
        # 「xとyを食べる」という情報
        if UFT.same(x, y) or UFT.same(x, y + 2 * N):
            ans += 1
        else:
            UFT.unite(x, y + N)
            UFT.unite(x + N, y + 2 * N)
            UFT.unite(x + 2 * N, y)

print(ans)