【競プロ典型90問】012の解説（python）

13797 ワード
AtCoder Union-Find木 Python3 AtCoder テキストリンク

概要

競プロ典型90問の解説記事です。
解説の画像を見ても分からない（理解力が足りない）ことが多々あったので、後で解き直した際に確認できるようまとめました。

問題012-Score Sum Queries

問題概要

白で塗られたH×Wのマス目に対して、以下の２パターンのいずれかのQ個のクエリを処理する。
1. 白いマス(r, c)を赤で塗る
2. (r1, c1), (r2, c2)の２点が与えられた時、以下の条件をどちらも満たしているか判定する
- ２点が共に赤く塗られている
- 2点が赤いマスを辿っていけば、連結している

解き方

この問題は、Union-Findを知っているか知らないかの問題です。
一応、簡単に説明すると、Union-Findとは、グループ分けの管理を非常に高速に実装できるデータ構造で、
ある集合のグループ分けや、同じグループに属しているかなどを高速に判定することができます。
Union-Findについての詳細は、ぜひ以下のスライドなどをご参照ください。


引用元：Union find(素集合データ構造) from Atcoder Inc.

Union-Findを用いると、各クエリの処理は以下のように、考えられます。
1. 白いマス(r, c)を赤で塗る
→ 白いマスを赤で塗り、隣接しているマスに赤いマスがあれば、連結させる。

２. (r1, c1), (r2, c2)の２点が与えられた時、赤いマスを辿っていけば、２点が連結している
→ 同じグループに属しているかを判定する

ここまで、手順が分かったので、以下で実際に実装していきたいと思います。

引用元：競プロ典型90問 Github

実装

answer.py

# 入力の受け取り
H, W = map(int, input().split())
Q = int(input())

# 二次元配列を一次元化する関数
def flatten(h, w):
  return h+(H+1)*w

# 赤塗りかどうかを判定する配列。配列の要素数をnに格納。
M = [False]*((H+1)*W)
n = len(M)

# 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 same(self, x, y):
    return self.find(x) == self.find(y)

# 要素数nを入れて、Union Findのインスタンスを作成。
uf = UnionFind(n)


for _ in range(Q):
  q = list(map(int, input().split()))

# ti = 1の時、[r, c]を赤く塗り、隣接マスを探索、もし赤なら、グループを結合する。
  if q[0] == 1:
    r, c = q[1]-1, q[2]-1
    x = flatten(r, c)
    M[x] = True
    for i, j in [[0, 1], [0, -1], [1, 0], [-1, 0]]:
      r_dash, c_dash = r+i, c+j
      y = flatten(r_dash, c_dash)
      if 0 <= r_dash < H and 0 <= c_dash < W:
        if M[y]:
          uf.union(x, y)
# ti = 2の時、[r1, c1], [r2, c2]が共に赤なら、同じグループに属するか判定。
  else:
    x = flatten(q[1]-1, q[2]-1)
    y = flatten(q[3]-1, q[4]-1)
    if M[x] and M[y] and uf.same(x, y):
      print("Yes")
    else:
      print("No")

最後に

問題の解説を読んでも他の人のコードを見てもさっぱり分からないという方の
力に少しでもなれれば幸いです。
コードの改善点やその他、ご意見、ご質問などあれば、気軽にコメントください。

ここまでお読みいただきありがとうございました。