Pythonで木構造内の特別なノードを見つけるプログラム

ここでは、「tree」という2次元リスト（n分木を表す）と、「color」という値のリストが与えられているとします。木は隣接リスト形式で表現されており、その根（ルート）は tree[0] です。

ノードの特徴

i番目のノードは以下の特徴を持ちます。

• tree[i]：そのノードの子ノードと親ノードの情報
• color[i]：そのノードの色

「特別なノード」とは

あるノードNを根とする部分木に含まれるすべてのノードの色が一意（重複なし）である場合、そのノードNを「特別（special）」なノードと呼びます。この定義に基づき、与えられた木の中に特別なノードがいくつあるかを求めるのが本問題の目的です。

入力例

tree = [
    [1,2],
    [0],
    [0,3],
    [2]
]
colors = [1, 2, 1, 1]

この場合、出力は 2 になります。

解法のアプローチ

この問題を解くために、以下の手順でDFS（深さ優先探索）を行います。

1. 結果を格納する変数 result を 0 で初期化します。
2. dfs(0, -1) を呼び出して、根から探索を開始します。
3. 最終的に result の値を返します。

check_intersection() 関数

引数として colors と child_colors を受け取ります。2つの集合間に共通する色が存在するかどうかを判定します。

• colors の長さが child_colors より短い場合：
  colors 内の各要素 c について、c が child_colors に存在すれば True を返します。
• それ以外の場合：
  child_colors 内の各要素 c について、c が colors に存在すれば True を返します。

小さい方の集合を走査することで、計算効率を高めています。

dfs() 関数

引数として node（現在のノード）と prev（直前のノード）を受け取ります。

• colors := {color[node]} として、現在のノードの色で集合を初期化します。
• tree[node] 内の各子ノード child について処理を行います。
  • child が prev と異なる場合（親への逆流を防ぐ）：
    • child_colors := dfs(child, node) で再帰的に子の部分木の色集合を取得します。
    • colors と child_colors がどちらも空でない場合：
      • check_intersection() が True を返した場合（色の重複あり）：colors := None とします。
      • 重複がない場合は、小さい方の集合に大きい方をマージして色集合を統合します。
    • どちらかが空の場合も colors := None とします。
• 最後に、colors が空でなければ result を1増やします（このノードは特別です）。
• colors を返します。

実装例

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

import collections
class Solution:
    def solve(self, tree, color):
        self.result = 0
        def dfs(node, prev):
            colors = {color[node]}
            for child in tree[node]:
                if child != prev:
                    child_colors = dfs(child, node)
                    if colors and child_colors:
                        if self.check_intersection(colors, child_colors):
                            colors = None
                        else:
                            if len(colors) < len(child_colors):
                                child_colors |= colors
                                colors = child_colors
                            else:
                                colors |= child_colors
                    else:
                        colors = None
            if colors:
                self.result += 1
            return colors
        dfs(0, -1)
        return self.result
    def check_intersection(self, colors, child_colors):
        if len(colors) < len(child_colors):
            for c in colors:
                if c in child_colors:
                    return True
        else:
            for c in child_colors:
                if c in colors:
                    return True
ob = Solution()
print(ob.solve( [
    [1,2],
    [0],
    [0,3],
    [2]
], [1, 2, 1, 1]))

入力

[
    [1,2],
    [0],
    [0,3],
    [2]
], [1, 2, 1, 1]

出力

2

まとめ

このアルゴリズムは、DFSを使って各ノードの部分木における色の集合をボトムアップに集約していきます。途中で色の重複が検出された時点で、そのノードは特別ではないと判断されます。計算量は木のノード数と色の管理方法に依存しますが、集合演算を活用することで効率的に判定できます。