典型90 012 – Red Painting（★4）｜かえるのほんだな

はじめに

競プロの精進のために学びがあった問題についてまとめています。目標は AtCoder 水色！

今のレートは茶色です。

問題

解説スライドはこちら。

考察

条件式のうち、「マス (rai,cai)(rai,cai) からマス (rbi,cbi)(rbi,cbi) まで赤色マス上を上下左右に移動することで辿り着ける」については、幅優先探索で実装することは可能ですが、計算量はO(N^2)になるため、今回のH、Wの最大値はそれぞれ2000なので、最悪計算量の想定の場合は耐えられません。

そこで、「マス (rai,cai)(rai,cai) からマス (rbi,cbi)(rbi,cbi) まで赤色マス上を上下左右に移動することで辿り着ける」という判定を高速化できる方法について考えます。

このような、「AからBにたどり着けるか」という問いはすなわち、「AとBが同じグループに存在するか」という判定問題に置き換えて、Union-Find木で効率的に探索することが可能です。

Union-Findについて

Union-findは、グループ分けを管理するためのデータ構造で、以下の2つの処理を高速に処理することができます。

ある要素とある要素が同一のグループに存在するかを調べる
ある要素とある要素を含むグループを、1つのグループに併合する。

そこで今回は、上下左右の移動でたどりつける連続した赤マスを1つのグループとしてとらえて、Union-find木を構築していきます。

Union-findを実装する場合にポイントになるのが、グループを根付き木で表現する点です。

ある数が同じグループに存在しているかについては、グループの根が一致するかで判断します。また、同じグループに存在する場合は、すべての子要素を直接根に繋ぎかえます。

子要素を繋ぎかえる際は、子要素数が小さい木の各要素を、子要素が大きい木に繋ぎかえる方が計算量が小さくてすみます。

というわけで、Union-findについてPythonで実装してみます。

class Union_find:

    def __init__(self, N):
        self.par = [-1] * (N+1)
        self.siz = [1] * (N+1)


    def root(self, x):
        if self.par[x] == -1:
            # root
            return x

        else:
            self.par[x] = self.root(self.par[x])
            return self.par[x]
        
    
    def issame(self, x, y):
        return self.root(x) == self.root(y)


    def unite(self, x, y):
        x = self.root(x)
        y = self.root(y)

        if x == y:
            return False
        
        else:
            if self.siz[x] < self.siz[y]:
                x, y = y, x
            
            self.par[y] = x;
            self.siz[x] += self.siz[y]

            return True


    def get_size(self, x):
        return self.siz[self.root(x)]


def main():
    N = 7
    uni = Union_find(100)
    uni.unite(1,2)
    uni.unite(2, 3)
    uni.unite(5, 6)

    print(uni.issame(1, 3))
    print(uni.issame(2, 5))

    uni.unite(1, 6)
    print(uni.issame(2, 5))


if __name__ == "__main__":
    main()

これをつかって、今回の問題についても解いてみました。

コード

class Union_find:
    
    def __init__(self, N):
        self.par = [-1] * (N+1)
        self.siz = [1] * (N+1)


    def root(self, x):
        if self.par[x] == -1:
            # root
            return x

        else:
            self.par[x] = self.root(self.par[x])
            return self.par[x]
        
    
    def issame(self, x, y):
        return self.root(x) == self.root(y)


    def unite(self, x, y):
        x = self.root(x)
        y = self.root(y)

        if x == y:
            return False
        
        else:
            if self.siz[x] < self.siz[y]:
                x, y = y, x
            
            self.par[y] = x;
            self.siz[x] += self.siz[y]

            return True


    def get_size(self, x):
        return self.siz[self.root(x)]

H, W = map(int,input().split())
Q = int(input())
G = Union_find(H*W+1)
M = [-1] * (H*W+1)

def check(n, x):
    if M[x] == 1 and M[n] == 1:
        G.unite(n, x)
    return

for i in range(Q):
    A = list(map(int,input().split()))
    if A[0] == 1:
        n = W * (A[1] - 1) + (A[2] - 1)
        M[n] = 1

        if n - W >= 0:
            check(n, n - W)
        if n + W < H*W:
            check(n, n + W)
        if n % W != 0 and n - 1 >= 0:
            check(n, n - 1)
        if (n + 1) % W != 0 and n + 1 < H*W:
            check(n, n + 1)

    else:
        n1 = W * (A[1] - 1) + (A[2] - 1)
        n2 = W * (A[3] - 1) + (A[4] - 1)

        if G.issame(n1, n2) and M[n1] == 1 and M[n2] == 1:
            print("Yes")

        else:
            print("No")