Day 090-Q1 — Segment Tree Merging（小到大マージ・部分木の最頻色集計）

問題

$N$ 頂点の根付き木（根は頂点1）が与えられる。頂点 $i$ には色 $c_i$（$1 \le c_i \le N$）が塗られている。各頂点 $v$ について、$v$ を根とする部分木に含まれる頂点の色の中で最も出現回数が多い色をすべて合計した値を求め、頂点 $1$ から $N$ まで順に出力せよ（最多出現回数の色が複数あればそれらの色番号をすべて足す）。

入出力例

4
1 2 1 3
1 2
1 3
1 4

1 2 1 3

5
1 1 2 2 3
1 2
2 3
2 4
1 5

3 2 2 2 3

例2: 頂点1の部分木の色は $[1,1,2,2,3]$。色1・色2がともに2回でタイのため答えは $1+2=3$。

概念図: 子の値域セグ木を親へマージ

ヒント

def merge(a, b, lo, hi):
    if a == 0: return b
    if b == 0: return a
    if lo == hi:
        cnt[a] += cnt[b]
        max_cnt[a] = cnt[a]; color_sum[a] = lo
        return a
    mid = (lo + hi) // 2
    left[a] = merge(left[a], left[b], lo, mid)
    right[a] = merge(right[a], right[b], mid+1, hi)
    # a の集約値を左右の子から pull() で再計算
    return a

模範解答

import sys

def main():
    input_data = sys.stdin.buffer.read().split()
    idx = 0
    n = int(input_data[idx]); idx += 1
    c = [int(input_data[idx + i]) for i in range(n)]; idx += n
    graph = [[] for _ in range(n + 1)]
    for _ in range(n - 1):
        u = int(input_data[idx]); idx += 1
        v = int(input_data[idx]); idx += 1
        graph[u].append(v)
        graph[v].append(u)

    limit = n * 40 + 10
    left = [0] * limit
    right = [0] * limit
    cnt = [0] * limit
    max_cnt = [0] * limit
    color_sum = [0] * limit
    node_count = 0

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

    def pull(node):
        l, r = left[node], right[node]
        lc, rc = max_cnt[l], max_cnt[r]
        if lc > rc:
            max_cnt[node] = lc; color_sum[node] = color_sum[l]
        elif rc > lc:
            max_cnt[node] = rc; color_sum[node] = color_sum[r]
        else:
            max_cnt[node] = lc; color_sum[node] = color_sum[l] + color_sum[r]
        cnt[node] = cnt[l] + cnt[r]

    def update(node, lo, hi, pos):
        if node == 0:
            node = new_node()
        if lo == hi:
            cnt[node] += 1
            max_cnt[node] = cnt[node]
            color_sum[node] = lo
            return node
        mid = (lo + hi) // 2
        if pos <= mid:
            left[node] = update(left[node], lo, mid, pos)
        else:
            right[node] = update(right[node], mid + 1, hi, pos)
        pull(node)
        return node

    def merge(a, b, lo, hi):
        if a == 0: return b
        if b == 0: return a
        if lo == hi:
            cnt[a] += cnt[b]
            max_cnt[a] = cnt[a]; color_sum[a] = lo
            return a
        mid = (lo + hi) // 2
        left[a] = merge(left[a], left[b], lo, mid)
        right[a] = merge(right[a], right[b], mid + 1, hi)
        pull(a)
        return a

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

    root = [0] * (n + 1)
    ans = [0] * (n + 1)
    for v in reversed(order):
        root[v] = update(root[v], 1, n, c[v - 1])
        for u in graph[v]:
            if u != parent[v]:
                root[v] = merge(root[v], root[u], 1, n)
        ans[v] = color_sum[root[v]]

    sys.stdout.write(' '.join(map(str, ans[1:])) + '\n')

main()

計算量: ならし $O(N \log N)$（各点更新が $O(\log N)$ ノードを生成し、マージは片方が空なら即終了するため総破棄数が生成数を超えない）。

Step-by-Step 解説

Step 1: 値域セグメント木を頂点ごとに動的生成

配列を毎回 $N$ サイズ確保すると $O(N^2)$ 空間になるため、必要なノードだけ動的に確保するポインタベースの実装にする。

Step 2: 反復DFSで処理順序を作る

再帰DFSはパス状の木で最悪 $O(N)$ 深度になりスタックオーバーフローの危険がある。スタックを使った反復DFSなら、頂点 $v$ の部分木全体が `order` 上で $v$ の直後に連続して現れるため、`reversed(order)` が正しい葉→根の処理順序になる。

Step 3: 自分の色を1点更新してから子をマージ

Step 4: マージのならし計算量

`merge` は毎回片方の入力ノードを吸収して破棄する形になる。全体で生成されるノード数は $O(N\log N)$ であり、マージで消費されるノード数もこれを超えない。

よくあるミス

次のステップ

• 発展問題: 部分木内で $k$ 番目に大きい出現回数の色を求めるクエリに拡張する（セグ木上二分探索）
• 発展問題: DSU on Tree（軽重分解）で同じ問題を解き実装量を比較する

自己評価