Pythonで同じラベルを持つサブツリー内のノード数を求めるプログラム

ここでは、n個のノードからなる根付きの一般木を考えます。ノードには0からn-1までの番号が振られており、各ノードには小文字の英字ラベルが割り当てられています。ラベルは配列labelsとして与えられ(labels[i]がi番目のノードのラベル)、木は辺リストで表現されます。各辺eは[u, v]という形式で、uが親、vが子であることを意味します。

求めたいのは、サイズnの配列Aです。A[i]には「i番目のノードと同じラベルを持つ、そのサブツリー内のノードの総数」を格納します。

例えば、入力が次のような場合を考えてみましょう。

n = 5、label = "ccaca" のとき、出力は [3, 2, 1, 1, 1] となります。これは、ルート(ノード0)には同じラベルを持つ子孫が3つ存在し、ノード1には2つ、それ以外のノードは自分自身だけが同じラベルを持つためです。

解法のアプローチ

この問題を解くために、以下の手順に従います。

• E := 与えられた辺リストからグラフ(隣接リスト)を作成する
• N := 各ノード番号と対応するラベルを格納したマップを作成する
• R := サイズnのリストを作成し、すべて0で初期化する
• 関数r()を定義する。引数はni(現在処理中のノード)
• C := ラベルの出現頻度を保持するマップ(Counter)
• E[ni]内の各辺eについて:
  • E[e]からniを削除する(親方向への逆流を防止)
  • r(e)の結果をCに反映する
• CにN[ni](自身のラベル)を反映する
• R[ni] := C[N[ni]] を設定する
• Cを返す
• メイン処理からr(0)を呼び出す
• Rを返す

それでは、以下の実装例を見ながら理解を深めていきましょう。

実装例(Python)

from collections import defaultdict, Counter

def solve(n, edges, labels):
    E = defaultdict(set)
    for f, t in edges:
        E[f].add(t)
        E[t].add(f)
    N = {i: e for i, e in enumerate(labels)}
    R = [0] * n

    def r(ni):
        C = Counter()
        for e in E[ni]:
            E[e].remove(ni)
            C.update(r(e))
        C.update((N[ni]))
        R[ni] = C[N[ni]]
        return C

    r(0)
    return R

n = 5
edges = [[0,1],[0,2],[1,3],[0,4]]
labels = "ccaca"
print(solve(n, edges, labels))

入力

5, [[0,1],[0,2],[1,3],[0,4]], "ccaca"

出力

[3, 2, 1, 1, 1]

アルゴリズムのポイント

このソリューションでは深さ優先探索(DFS)を活用しています。再帰関数r()は各ノードを訪問し、そのサブツリー全体のラベル出現頻度をCounterとして集計します。子ノードの結果を順にマージし、最後に自身のラベルを加算することで、各ノードのサブツリー内における同ラベルのノード数を効率的に計算できます。

また、辺リストからは一旦無向グラフを構築し、DFSの過程で隣接リストから親ノードを削除していくことで、木を子方向へのみ辿るように制御している点も重要なテクニックです。これにより、循環を避けながら全ノードを正確に一度ずつ処理できます。