何の話かと言うと
上記の問題をネタに、（木構造とは限らない）単純無向グラフの基本的な取り扱いを説明します。
単純無向グラフは、多重リンクや自分から自分へのリンクを含まない無向グラフです。一般には、複数の連結成分に別れます。
データの持ち方
基本的には、各ノードに対して、リンクされたノードのリストを用意すればOKです。
links = [[] for _ in range(N+1)] for _ in range(M): a, b = list(map(int, f.readline().split())) links[a].append(b) links[b].append(a)
連結成分に対する処理
次のようにキューを用いると、あるノード n に対して、そこからすべての連結成分を辿ることができます。
flag = [None] * (N+1) # 処理済みのノードを示す q = deque() flag[n] = True q.append(n) while q: i = q.pop() # ここでノード i についての処理をする for j in links[i]: if flag[j]: continue flag[j] = True q.append(j)
木構造の場合は、下記の記事で説明したように、親ノードの情報を与えることで「上から下にたどる」という順序を保証しました。
今回は木構造とは限らないので、処理済みのノードを示すフラグを利用して、ループを避けています。木構造の場合も同じ手法を使うことができます。
連結部分ごとに個別に処理
連結とは限らないグラフに対して、連結部分ごとに処理をしたい場合は、次の様にすべてのノードについてループすればOKです。
flag = [None] * (N+1) # 処理済みのノードを示す for n in range(1, N+1): if flag[n]: continue q = deque() flag[n] = True q.append(n) while q: i = q.pop() # ここでノード i についての処理をする for j in links[i]: if flag[j]: continue flag[j] = True q.append(j)
フラグのチェックによって、処理済みのノードはスキップされるので、同じ連結成分を2回処理することはありません。
問題の解法
で、最後に冒頭の問題の解法ですが、連結部分ごとに場合の数を計算して、それらの積で全体の答えになることから、上記のパターンが使えそうだと気がつきます。
それぞれの連結部分について、ノードを深さ優先探索でたどりながら、その時点で可能な配色を網羅していきます。
たとえば、
2 --- 1 --- 3 |-----------|
というループ状のグラフを 1, 2, 3 の順に処理する場合、
1. 可能な配色は、[R, None, None] （対称性から、G, B の場合はスキップして最後に3倍する。None は未定の意味）
2.可能な配色は、[R, G, None], [R, B, None]
3. 可能な配色は、[R, G, B], [R, B, G]
という様に順番に決まっていきます。あるノードを処理する際に、（深さ優先探索をしているので）そこからリンクされたノードの1つは確実に処理済みなので、追加する色の候補は、2種類以下になります。これにより、組み合わせの数を制限して、組合せ爆発を防いでいます。
具体的な実装としては、各ステップにおいて、その時点で可能性のある配色全体のリスト（candidate = [[R, G, None], [R, B, None]] みたいなやつ）をアップデートしていく感じになります。この配色リストをアップデートする処理を上記のパターンの
# ここでノード i についての処理をする
という部分に当てはめればOKです。
ただし、この手のアップデートパターン（一種のDPですね）では、最初のノードに対する処理だけは別枠になるので、そのまま当てはめると場合分けの記述がごちゃごちゃします。そこで、まずは、深さ優先探索でグラフを辿る際のノード順のリストを作ります。
tree_order = []
q = deque()
q.append(n)
flag[n] = True
while q:
i = q.pop()
tree_order.append(i)
for j in links[i]:
if flag[j]:
continue
flag[j] = True
q.append(j)
この後、このリスト順にループを回せば、最初の要素 tree_order[0] に対する処理をきれいに分けて書くことができます。
完成版はこんな感じ。
import sys, copy from collections import defaultdict, deque def main(f): N, M = list(map(int, f.readline().split())) links = [[] for _ in range(N+1)] for _ in range(M): a, b = list(map(int, f.readline().split())) links[a].append(b) links[b].append(a) flag = [None] * (N+1) result = 1 for n in range(1, N+1): if flag[n]: continue tree_order = [] q = deque() q.append(n) flag[n] = True while q: i = q.pop() tree_order.append(i) for j in links[i]: if flag[j]: continue flag[j] = True q.append(j) candidates = [] i = tree_order[0] key = [None] * (N+1) key[i] = 0 candidates.append(key) for i in tree_order[1:]: candidates_n = [] for col_list in candidates: for c in range(3): ok = True for j in links[i]: if col_list[j] == c: ok = False if ok: col_list[i] = c candidates_n.append(copy.copy(col_list)) candidates = candidates_n result *= len(candidates) * 3 print(result) main(sys.stdin)
