【競プロ典型90問】012 - Red Painting（★4）

概要・感想

H x Wのマス目の上でUnionFindを行うという問題。問題の解法すぐに浮かんだが、ソラでUnionFindを書こうとしたら思いのほかバグを発生させてしまい、時間がかかった。

具体的には、UnionFind.merge()内のif a==b:部分を書き忘れていてどこが間違っているのかわからなかった。

import sys
sys.setrecursionlimit(10**7)

class UnionFind:
    def __init__(self, n):
        self.p = [-1] * n
    def parent(self, a):
        if self.p[a] < 0:
            return a
        else:
            self.p[a] = self.parent(self.p[a])
            return self.p[a]
    def is_same(self, a, b):
        a = self.parent(a)
        b = self.parent(b)
        return a == b
    def merge(self, a, b):
        a = self.parent(a)
        b = self.parent(b)
        if a == b:
            return
        if self.p[a] > self.p[b]:
            a, b = b, a
        self.p[a] += self.p[b]
        self.p[b] = a

h, w = map(int, input().split())
M = 2200
uf = UnionFind(M * M)
qn = int(input())

col = [0] * (M * M)
dd = [M, -M, 1, -1]


for _ in range(qn):
    q = list(map(int, input().split()))
    if q[0] == 1:
        x = q[1] * M + q[2]
        col[x] = 1
        for d in dd:
            if col[x+d] == 1:
                uf.merge(x, d + x)
    else:
        x = q[1] * M + q[2]
        y = q[3] * M + q[4]
        print('No' if not uf.is_same(x, y) or col[x] == 0 else 'Yes')