Pythonで二分木の隣接ノードが同じ色にならないように着色できるか判定するプログラム

各ノードの値がそのノードの色を表す二分木を考えます。木に含まれる色は最大で2色です。ここで、ノード同士の色を何度でも入れ替えられるとき、辺でつながれた隣接ノード同士が同じ色にならないような配置が可能かどうかを判定します。

たとえば、入力が次のような木だったとします。

この場合の出力は True です。色を入れ替えることで、次のようにすべての隣接ノードが異なる色になる状態を作れるからです。

解法のアプローチ

この問題は、次の手順で解くことができます。

• colors := 空のマップ(各色を持つノードの個数を記録)
• prop := 空のマップ(フラグごとのノード数を記録)
• dfs() 関数を定義する。引数は node と flag
• node が null の場合は何もせずに return
• colors[node の値] を 1 増やす
• prop[flag] を 1 増やす
• dfs(node の左の子、flag を反転した値)を呼び出す
• dfs(node の右の子、flag を反転した値)を呼び出す
• メイン処理では dfs(root, True) を実行し、colors と prop の値がすべて一致していれば True、そうでなければ False を返す

このアルゴリズムの鍵となるのは、木が二部グラフであるという性質です。DFSで隣接ノードに交互にフラグ(True / False)を割り当てることで、木全体を矛盾なく2つのグループに分割できます。色はノード間で自由に入れ替えられるため、各グループのサイズと各色のノード数がうまく対応づけられれば、「隣接ノードが必ず異なる色になる」という条件を満たす配置を実現できます。

それでは、理解を深めるために以下の実装を見てみましょう。

実装例

from collections import defaultdict
class TreeNode:
    def __init__(self, data, left=None, right=None):
        self.val = data
        self.left = left
        self.right = right

class Solution:
    def solve(self, root):
        colors = defaultdict(int)
        prop = defaultdict(int)

        def dfs(node, flag=True):
            if not node:
                return
            colors[node.val] += 1
            prop[flag] += 1
            dfs(node.left, not flag)
            dfs(node.right, not flag)

        dfs(root)
        return set(colors.values()) == set(prop.values())

ob = Solution()
root = TreeNode(2)
root.left = TreeNode(2)
root.right = TreeNode(2)
root.right.left = TreeNode(1)
root.right.right = TreeNode(1)
root.right.left.right = TreeNode(1)
print(ob.solve(root))

入力

root = TreeNode(2)
root.left = TreeNode(2)
root.right = TreeNode(2)
root.right.left = TreeNode(1)
root.right.right = TreeNode(1)
root.right.left.right = TreeNode(1)

出力

True