Day 117-Q5 | DSU on Tree（small-to-largeマージ）

問題

$N$ 頂点の木が頂点 $1$ を根として与えられる。各頂点 $v$ には色 $C_v$（$1\le C_v\le N$）が塗られている。

全ての頂点 $v$ について、$v$ を根とする部分木に含まれる相異なる色の種類数を求めよ。

入力形式

N
a_1 b_1
a_2 b_2
:
a_{N-1} b_{N-1}
C_1 C_2 ... C_N

制約

$2 \le N \le 2\times10^5$

$1 \le a_i, b_i \le N$（木を成す辺）

$1 \le C_v \le N$

入出力例

入力例1

出力例1

3 2 1 1 1

頂点1の部分木$\{1,2,3,4,5\}$、色$\{1,2,1,3,2\}$で相異なる色は3種類。頂点2の部分木$\{2,4,5\}$、色$\{2,3,2\}$で2種類。頂点3,4,5は葉なので1種類ずつ。

概念図: heavy child（重い子）だけ情報を持ち越す

ヒント（段階的開示）

ヒント1: 方向性

各頂点について愚直に部分木を辿って集合を作ると、1回のクエリに $O(N)$ かかり全体で $O(N^2)$ になってしまう。木上のクエリを高速化する定番テクニックとして、「軽い部分木は毎回作り直して捨てる、重い部分木（最大の子）だけは情報を使い回す」という考え方がある。

ヒント2: アプローチ

これは small-to-large マージ（別名 DSU on Tree）と呼ばれる手法。各頂点でまず「最も部分木サイズが大きい子（heavy child）」を1つ選んでおく。DFSで、まず軽い子たちを再帰的に処理してその情報は使い終わったら破棄する。次にheavy childを再帰的に処理しその情報は保持したままにする。最後に軽い子たちの情報をheavy childの情報に足し込み、自分自身の色も加えて答えを確定させる。これにより全体 $O(N\log N)$ で計算できる。

ヒント3: 誘導（コード骨格）

# 1. 各頂点の部分木サイズを求め、最大の子を heavy child として記録する
# 2. dfs(v, keep):
#      軽い子を dfs(..., keep=False) で再帰（終わったら情報は破棄済み）
#      heavy child があれば dfs(heavy, keep=True) で再帰（色カウントが残る）
#      軽い子の部分木の色を、色カウント辞書に1つずつ加算する
#      自分の色も加算する
#      ans[v] = 現在の「相異なる色の種類数」
#      keep が False なら、加算した分を全部差し引いて元に戻す

模範解答 (Python)

import sys
sys.setrecursionlimit(300000)


def solve():
    data = sys.stdin.buffer.read().split()
    idx = 0
    n = int(data[idx]); idx += 1
    adj = [[] for _ in range(n + 1)]
    for _ in range(n - 1):
        a = int(data[idx]); b = int(data[idx + 1]); idx += 2
        adj[a].append(b)
        adj[b].append(a)
    color = [0] + [int(data[idx + i]) for i in range(n)]
    idx += n

    # 部分木サイズと heavy child を求める（親付きBFS順で）
    parent = [0] * (n + 1)
    order = []
    size = [1] * (n + 1)
    visited = [False] * (n + 1)
    stack = [1]
    visited[1] = True
    while stack:
        v = stack.pop()
        order.append(v)
        for to in adj[v]:
            if not visited[to]:
                visited[to] = True
                parent[to] = v
                stack.append(to)
    for v in reversed(order):
        for to in adj[v]:
            if to != parent[v]:
                size[v] += size[to]

    heavy = [0] * (n + 1)
    for v in range(1, n + 1):
        best = -1
        for to in adj[v]:
            if to != parent[v] and size[to] > best:
                best = size[to]
                heavy[v] = to

    cnt = {}
    distinct = [0]

    def add_light(v, par, delta, skip):
        c = color[v]
        cnt[c] = cnt.get(c, 0) + delta
        if delta == 1 and cnt[c] == 1:
            distinct[0] += 1
        if delta == -1 and cnt[c] == 0:
            distinct[0] -= 1
        for to in adj[v]:
            if to != par and to != skip:
                add_light(to, v, delta, skip)

    ans = [0] * (n + 1)

    def dfs(v, par, keep):
        for to in adj[v]:
            if to != par and to != heavy[v]:
                dfs(to, v, False)
        if heavy[v]:
            dfs(heavy[v], v, True)
        for to in adj[v]:
            if to != par and to != heavy[v]:
                add_light(to, v, 1, -1)
        c = color[v]
        cnt[c] = cnt.get(c, 0) + 1
        if cnt[c] == 1:
            distinct[0] += 1
        ans[v] = distinct[0]
        if not keep:
            add_light(v, par, -1, -1)

    dfs(1, 0, False)
    print(*ans[1:n + 1])


solve()

計算量: small-to-largeマージにより、各頂点が「軽い部分木として加算・削除される」回数の合計は $O(N\log N)$ に抑えられる。入力例1、および300ケースのランダム木に対する愚直DFS集計との比較で正しさを検証済み。

Step-by-Step 解説

1部分木サイズとheavy childの事前計算
根からのBFS/DFSで各頂点の部分木サイズを求め、各頂点について「子の中で最も部分木サイズが大きいもの」をheavy childとして記録しておく。

2軽い子を先に処理して捨てる
dfs(v, keep=False)で軽い子を再帰的に処理する。keep=Falseで呼ばれた頂点は、自分の処理が終わった後に自分が加算した情報を全て取り消す。

3heavy childだけ情報を持ち越す
heavy childはdfs(heavy, keep=True)で呼び出す。色カウントの辞書cntが破棄されずそのまま残るので、軽い子たちの色をadd_lightで足し込むだけで部分木全体の色集計が完成する。

4計算量が$O(N\log N)$になる理由
ある頂点が「軽い子」として加算される回数は、根に向かって「軽い辺」を何回通るかに等しい。軽い辺を1回通るたびに部分木サイズは少なくとも半分以下になるため、軽い辺の連続回数は高々$O(\log N)$回。

5実装上の注意
add_lightはheavy childを常に飛ばして呼び出す必要がある。skipを渡し忘れると二重に数えてしまうバグになる。

よくあるミス

ミス	原因	正しい書き方
全ての子を同列に扱い、毎回情報を作り直してしまう	heavy childの概念自体を実装に反映し忘れる	部分木サイズ最大の子だけは情報を保持したまま次に活かす（`keep=True`）
`add_light`でheavy child方向にも再帰してしまう	`skip`パラメータを正しく使っていない	heavy childは既に`cnt`に反映済みなので明示的に除外する
`keep=False`の頂点で情報を取り消し忘れる	「使い終わったら破棄する」手順を実装し忘れる	関数末尾で`keep`が偽の場合は必ず加算した分を全て差し引く
部分木サイズを通常のDFS順で計算し、子の情報が揃う前に確定させてしまう	帰りがけ順（postorder）で計算する必要性を見落とす	BFSで訪問順を記録し、その逆順（葉から根へ）でサイズを積み上げる

次のステップ

発展: セグメント木マージ（Segment Tree Merging）でも同じ問題を解き、DSU on Treeとの実装・計算量の違いを比較する
次回予告: 未定（Master Levelローテーション継続）

自己評価

完全理解おおむね理解要復習

Day 117-Q5 — DSU on Tree（small-to-largeマージによる部分木クエリ）

問題