Day 139-Q5 | Segment Tree Merging

問題

$N$ 頂点の根付き木（根は頂点 $1$）が与えられる。各頂点 $i$ には色 $c_i$（$1\le c_i\le N$）が塗られている。

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

この問題は「small-to-large マージ（DSU on Tree）」で解くのが定番だが、本問ではあえて Segment Tree Merging（セグメント木マージ） で解く。各頂点に「色の値域 $[1,N]$ 上の動的セグメント木」を1本持たせ、子から親へボトムアップにマージしていく。2本のセグメント木をマージする際、両方に実体があるノード同士だけ再帰的に合成し、片方が存在しないノードはそのまま相手を採用することで、マージ全体の総コストを償却 $O(N\log N)$ に抑えられる。

入力形式

N
c_1 c_2 ... c_N
u_1 v_1
:
u_{N-1} v_{N-1}

制約

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

$1 \le c_i \le N$

与えられる辺は頂点1を根とする木を成す

入出力例

入力例1

出力例1

3 1 3 1 1

概念図

ヒント（段階的開示）

ヒント1（方向性）

「部分木に含まれる相異なる値の個数」は、DSU on Tree（dict/set を small-to-large でマージ）で解くのが標準的だが、その場合ハッシュ集合のマージコストが実装依存で振れやすい。代わりに「値の存在」をセグメント木の1リーフ＝1点として表現し、木構造そのものをマージすることで、同じ計算量 $O(N\log N)$ をより機械的な操作で達成できないか、と考えるのが出発点。

ヒント2（アプローチ）

各頂点 $v$ に対して値域 $[1,N]$ 上の動的セグメント木を1本作る。まず自分の色 $c_v$ の位置に1点更新する。その後、子（すでに処理済み）のセグメント木を自分のセグメント木にマージする。マージ関数は「どちらか一方が存在しないなら、もう一方をそのまま返す」「両方とも葉なら、どちらか一方をそのまま返す（値を加算しない）」「両方とも内部ノードなら左右の子を再帰的にマージする」という規則で定義する。

ヒント3（誘導）

def merge(a, b):
    if a == 0:
        return b
    if b == 0:
        return a
    if left[a] == 0 and right[a] == 0 and left[b] == 0 and right[b] == 0:
        return a  # 同じ色の位置。加算せずどちらかを採用
    left[a] = merge(left[a], left[b])
    right[a] = merge(right[a], right[b])
    val[a] = val[left[a]] + val[right[a]]
    return a

木を「親より子を先に処理する」順序で辿り、各頂点のセグメント木を作ってから親のセグメント木にマージしていけばよい。

模範解答 (Python)

import sys


def main():
    data = sys.stdin.buffer.read().split()
    idx = 0
    n = int(data[idx]); idx += 1
    c = [0] * (n + 1)
    for i in range(1, n + 1):
        c[i] = int(data[idx]); idx += 1

    adj = [[] for _ in range(n + 1)]
    for _ in range(n - 1):
        u = int(data[idx]); idx += 1
        v = int(data[idx]); idx += 1
        adj[u].append(v)
        adj[v].append(u)

    MAXV = n
    max_nodes = (n + 5) * 22 + 50
    left = [0] * max_nodes
    right = [0] * max_nodes
    val = [0] * max_nodes
    node_cnt = 0

    def new_node():
        nonlocal node_cnt
        node_cnt += 1
        return node_cnt

    def update(node, l, r, pos):
        if node == 0:
            node = new_node()
        if l == r:
            val[node] = 1
            return node
        mid = (l + r) // 2
        if pos <= mid:
            left[node] = update(left[node], l, mid, pos)
        else:
            right[node] = update(right[node], mid + 1, r, pos)
        val[node] = val[left[node]] + val[right[node]]
        return node

    def merge(a, b):
        if a == 0:
            return b
        if b == 0:
            return a
        if left[a] == 0 and right[a] == 0 and left[b] == 0 and right[b] == 0:
            return a
        left[a] = merge(left[a], left[b])
        right[a] = merge(right[a], right[b])
        val[a] = val[left[a]] + val[right[a]]
        return a

    root = [0] * (n + 1)
    parent = [0] * (n + 1)
    visited = [False] * (n + 1)
    order = []
    stack = [1]
    visited[1] = True
    while stack:
        u = stack.pop()
        order.append(u)
        for v in adj[u]:
            if not visited[v]:
                visited[v] = True
                parent[v] = u
                stack.append(v)

    # 各頂点に自分の色を1点更新
    for u in order:
        root[u] = update(root[u], 1, MAXV, c[u])

    # order は「親が子より先」に並んでいるので、逆順で処理すれば「子が親より先」になる
    for u in reversed(order):
        if parent[u] != 0:
            root[parent[u]] = merge(root[parent[u]], root[u])

    ans = [str(val[root[u]]) for u in range(1, n + 1)]
    print(' '.join(ans))


main()

Step-by-Step 解説

1なぜ「木構造そのもの」をマージできるのか
各頂点のセグメント木は「値域上のどこに色が存在するか」を表す木である。マージは集合としての和集合（OR）を取ることに相当する。ノードが「存在しない」側は情報を持たないので、もう一方をそのまま使えばよい。

2償却計算量の直感
マージ関数は「両方とも非0ノード」の場合にのみ再帰を続ける設計になっているため、再帰が深く進むのは2つの木が重なる部分だけである。木全体で作られるノードの総数は $O(N\log N)$ 個に抑えられ、マージは既存ノードを消費する操作なので新規ノードを一切生成しない。したがって全体で $O(N\log N)$ に収まる。

3update と merge の役割分担
update は「1つの頂点の色を、その頂点専用のセグメント木に反映する」処理（$O(\log N)$ のノード生成）。merge は「2つの完成した木を1つに統合する」処理で、既存ノードの再利用によって計算量を抑える。

4処理順序（子を先に処理する）
頂点1からのスタックDFSで得られる訪問順は「親が子より先に現れる」順序になっている。これを reversed() すれば「子が親より先に処理される」順序になり、親のセグメント木への安全なマージが可能になる。

5DSU on Tree との比較
DSU on Tree は「軽い子の集合を重い子の集合へ吸収させる」ことで償却 $O(N\log N)$ を達成する。Segment Tree Merging は「マージのたびに少なくとも1つのノードが消滅する」というポテンシャル論法で償却計算量を導く。どちらも同じ計算量オーダーだが、後者はセグメント木で表現できる情報であれば統計量を変えるだけで幅広く応用が利く。

よくあるミス

ミス	原因	正しい書き方
merge の中で毎回新規ノードを作ってしまう	通常の「永続セグメント木」の実装（コピーオンライトで新規ノードを作る）と混同する	本問の merge は既存ノード a を破壊的に更新して再利用する（新規ノードを作らない）ことで計算量を抑える設計にする
マージ時に両方とも葉の場合、値を加算してしまう（val[a] += val[b]）	「区間和クエリ」の一般的なマージ操作（値の合算）と、本問の「存在の統合（OR）」を混同する	同じ色の位置（両方とも葉）なら加算せずどちらか一方をそのまま採用する（重複カウントを避ける）
処理順序を間違え、親のセグメント木がまだ空の状態で子をマージしようとする	DFSの訪問順をそのまま使い、子を先に処理する順序になっていることを確認しない	スタックベースDFSで得た訪問順（親が子より先）を reversed() して使う
max_nodes（配列サイズ）を過小に見積もり、実行時に配列外アクセスが起きる	1点更新1回あたり O(log N) 個のノードを作ることを忘れ、N個分しか確保しない	ノード総数の上限は「更新回数×木の高さ」なので、N×(log2 N+2)程度の余裕を持たせて確保する
木を明示的スタックで辿らず再帰DFSで書き、大きなNでスタックオーバーフローを起こす	Pythonの再帰関数呼び出し回数の上限を意識していない	訪問順の構築は明示的スタックで行い、再帰は update/merge（深さO(log N)）にのみ使う

次のステップ

発展: 「部分木内の相異なる色の個数」ではなく「部分木内で最も出現回数が多い色とその出現回数」を、同じセグメント木マージの枠組みで求めるプログラムに拡張せよ。
発展2: 本問をDSU on Tree（dict/Counter を small-to-large マージ）でも実装し、$N$ を $2\times10^5$ 程度まで振って、Segment Tree Merging との実行時間・メモリ使用量を比較計測せよ。

自己評価

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

Day 139-Q5 — Segment Tree Merging（動的セグメント木のマージ技法）

問題