Atcoder 競プロ典型 90 問 012 - Red Painting

問題文は👆を確認してください

導入

atcoderで精進してますが、
どこかにアウトプットしたかったので、qiitaに書くことにしました。
ACしてから載せるようにしてますが、考え方が至っていない点があればご指摘頂けると幸いです。
また、誰かの参考になると嬉しいです。

使用しているアルゴリズム

• Union-Find（素集合データ構造）
• グリッドの2次元→1次元変換

解法のポイント

この問題は、グリッド上の赤マス同士が連結しているかどうかを、高速に判定する必要があります。
そこで、以下の工夫をしています：

🔹グリッドのUnion-Find化

• 各マスを「2次元→1次元」へ h*W + w で変換し、Union-Findで管理。
• 初期状態ではすべて白（未訪問）とし、赤マスに変えた際に上下左右の赤マスと連結（merge）します。

🔹Union-Find実装の詳細

• find(j)：親を再帰的にたどりつつ、経路圧縮も実施
• merge(i, j)：根っこ同士を比較して連結
• same(i, j)：2つのマスが同じ連結成分にいるか判定

🔹クエリ処理

• クエリ1（赤く塗る）：該当マスを赤（visited=True）にし、上下左右の赤マスがあればUnion
• クエリ2（赤マス間連結判定）：両マスが赤で、かつUnion-Findで同じ親を持っていれば Yes

ACサンプル

class UnionFind:
  def __init__(self,n):
    self.parent = [-1] * n
  
  def find(self,j):
    if self.parent[j] == -1:
      return j
    self.parent[j] = self.find(self.parent[j])
    return self.parent[j]
  
  def merge(self,i,j):
    i_root = self.find(i)
    j_root = self.find(j)
    if not i_root == j_root:
      self.parent[i_root] = j_root
  
  def same(self,i,j):
    i_root = self.find(i)
    j_root = self.find(j)
    return i_root == j_root

H,W = map(int,input().split())
N = int(input())
visited = [[False]*(W+2) for _ in range(H+2)]
uf = UnionFind(H*W)

def query1(h,w):
  visited[h][w] = True
  dhs = [1,-1,0,0]
  dws = [0,0,1,-1]
  for dh,dw in zip(dhs,dws):
    fh = h + dh
    fw = w + dw
    if fh < 0 or fh >= H or fw < 0 or fw >= W:
      continue
    if visited[fh][fw]:
      uf.merge(h*W+w, fh*W+fw)

def query2(ah,aw,bh,bw):
  if visited[ah][aw] and visited[bh][bw]:
    return uf.same(ah*W+aw,bh*W+bw)
  return False
  
for _ in range(N):
  q = list(map(int,input().split()))
  if q[0] == 1:
    query1(q[1]-1,q[2]-1)
  else:
    if query2(q[1]-1,q[2]-1,q[3]-1,q[4]-1):
      print("Yes")
    else:
      print("No")

リンク

ACサンプル集(アルゴリズム逆引き)