Official

E - 水路の整備 / Maintenance of Waterways Editorial by admin

Claude 4.6 Opus (Thinking)

概要

根付き木の各辺を、その辺の両端のどちらの集落に担当させるかを決める最適化問題です。各集落には許容量があり、超過すると追加コストが発生するため、木DPとグリーディを組み合わせて最小コストを求めます。

考察

重要な気づき

  1. 辺ごとの独立した選択: 各辺 \((i, P_i)\) について「親 \(P_i\) に割り当てる」か「子 \(i\) に割り当てる」かの二択を全辺分決定する問題です。

  2. 木DP の定式化: 各集落で「自分がいくつの辺を担当するか」を決める必要があり、部分木単位で最適解を構築できます。

  3. 限界コストの構造: 集落 \(i\) に割り当てる辺を1本ずつ増やすとき、\(C_i\) 本目まではペナルティ 0、\(C_i + 1\) 本目以降は1本あたり \(W_i\) のペナルティです。この限界コストが単調非減少であるため、貪欲法が適用可能です。

DP の状態定義

\(f[v][b]\) を「集落 \(v\) の部分木内で発生する追加コスト(ペナルティ)の最小値」と定義します: - \(b = 0\): 辺 \((v, P_v)\) を親 \(P_v\) が担当する場合 - \(b = 1\): 辺 \((v, P_v)\) を集落 \(v\) が担当する場合

アルゴリズム

基本方針

各ノード \(v\) の子 \(c\) について、辺 \((v, c)\) を: - 子 \(c\) に割り当てる → 子の部分木コストは \(f[c][1]\) - 親 \(v\) に割り当てる → 子の部分木コストは \(f[c][0]\)\(v\) の担当数 +1

貪欲な切り替え

  1. 初期状態: 全子辺を子側に割り当てた状態から出発(コスト \(= \sum_c f[c][1]\)
  2. 切り替えの利得: 子 \(c\) の辺を \(v\) に移すと、部分木コストが \(d = f[c][1] - f[c][0]\) だけ減少する
  3. 限界ペナルティ: \(v\) の現在の担当数が \(C_v\) 以上なら \(W_v\)、未満なら \(0\)
  4. 判定: 利得 \(d\) が限界ペナルティ \(mp\) を超える限り切り替えを実行

\(d\) の大きい子から順に処理し、\(d \leq mp\) になった時点で打ち切ります(\(d\) は降順ソート済み、\(mp\) は非減少なので最適)。

計算式

\[f[v][b] = \left(\sum_c f[c][1]\right) - \sum_{\text{切り替えた子}} (d - mp) + W_v \cdot \max(0, b - C_v)\]

最終的な答えは \((N-1) + f[0][0]\) です。

計算量

  • 時間計算量: \(O(N \log N)\)(各ノードで子のリストをソートする合計コスト)
  • 空間計算量: \(O(N)\)

実装のポイント

  • 根の特別処理: 根(集落1)には親辺がないため、\(b = 1\) の状態は存在しません。コード中では f[0][1] = 0 として continue しています。

  • BFS的な順序での処理: 再帰の代わりにスタックを用いてトポロジカル順序を求め、逆順(葉から根へ)に処理することでスタックオーバーフローを回避しています。

  • ソートによる貪欲: \(d = f[c][1] - f[c][0]\) を降順ソートすることで、利得の大きいものから優先的に割り当て、限界コストが上がった時点で打ち切る正当性が保証されます。

    ソースコード

import sys
from sys import setrecursionlimit

def main():
    input_data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(input_data[idx]); idx += 1
    children = [[] for _ in range(N)]
    for i in range(1, N):
        p = int(input_data[idx]) - 1; idx += 1
        children[p].append(i)
    C = [0] * N; W = [0] * N
    for i in range(N):
        C[i] = int(input_data[idx]); idx += 1
        W[i] = int(input_data[idx]); idx += 1

    f = [[0, 0] for _ in range(N)]
    order = []
    stack = [0]
    while stack:
        v = stack.pop()
        order.append(v)
        for c in children[v]:
            stack.append(c)
    for v in reversed(order):
        base = sum(f[c][1] for c in children[v])
        ds = sorted([f[c][1] - f[c][0] for c in children[v]], reverse=True)
        for b in range(2):
            if v == 0 and b == 1:
                f[v][1] = 0
                continue
            cost = base
            taken = b
            benefit = 0
            for d in ds:
                marginal_penalty = W[v] if taken >= C[v] else 0
                net = d - marginal_penalty
                if net > 0:
                    benefit += net
                    taken += 1
                else:
                    break
            cost -= benefit
            cost += W[v] * max(0, b - C[v]) if b > 0 else 0
            # Recalculate properly
            total = base
            taken = b
            for d in ds:
                mp = W[v] if taken >= C[v] else 0
                if d > mp:
                    total -= d; total += mp
                    taken += 1
                else:
                    break
            f[v][b] = total + W[v] * max(0, b - C[v])
    print((N - 1) + f[0][0])

main()

この解説は claude4.6opus-thinking によって生成されました。

posted:
last update: