Official

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

gemini-3.5-flash-thinking

概要

本問題は、木構造の道路網において、各辺(道路)の水路工事を隣接する2つの集落のどちらに割り当てるかを決定し、発生する追加コストの総和を最小化する問題です。 木構造の性質を活かし、葉から根に向かってボトムアップに最適な割り当てを決定していく木DP(動的計画法)を用いることで、効率的に解くことができます。


考察

1. 意思決定の構造化

各道路(辺)は、親集落と子集落を結んでいます。この道路の工事担当は「親が引き受ける」か「子が引き受ける」かの2通りしかありません。 したがって、各部分木において、「その部分木の根と親を結ぶ道路をどちらが担当するか」によって状態を分けて考えることができます。

具体的には、各集落 \(u\) について以下の2つの状態(最小コスト)を定義します。 - \(dp_0[u]\): 集落 \(u\) とその親 \(P_u\) を結ぶ道路を\(P_u\) が担当する場合の、集落 \(u\) の部分木における最小追加コスト。 - \(dp_1[u]\): 集落 \(u\) とその親 \(P_u\) を結ぶ道路を\(u\) が担当する場合の、集落 \(u\) の部分木における最小追加コスト。

2. 子集落からの遷移

集落 \(u\) の子集落の集合を \(Ch(u)\) とします。集落 \(u\) が、子集落との間の道路(計 \(k\) 本)のうち、ちょうど \(c\)を自身で担当すると仮定します(\(0 \leq c \leq k\))。

このとき、選ばれた \(c\) 個の子集落 \(v\) については、道路 \((u, v)\)\(u\) が担当するため、子 \(v\) から見ると「親が担当する」ことになり、コストは \(dp_0[v]\) となります。 選ばれなかった残り \(k - c\) 個の子集落 \(v\) については、道路 \((u, v)\)\(v\) が担当するため、コストは \(dp_1[v]\) となります。

したがって、集落 \(u\) が担当する子集落の集合を \(S \subseteq Ch(u)\)(ただし \(|S| = c\))とすると、子集落から得られるコストの総和は以下のようになります。 $\( \sum_{v \in S} dp_0[v] + \sum_{v \notin S} dp_1[v] \)$

これを変形すると、 $\( \sum_{v \in Ch(u)} dp_1[v] + \sum_{v \in S} (dp_0[v] - dp_1[v]) \)$ となります。

3. 貪欲法による最適化

任意の \(c\) に対して上記の値を最小化するには、差分 \(D[v] = dp_0[v] - dp_1[v]\)小さいものから順に \(c\)を選んで \(S\) に入れればよいです。 あらかじめ \(D[v]\) を昇順にソートしておくことで、すべての \(c \in [0, k]\) に対する最小コストを \(O(k \log k)\) で求めることができます。

集落 \(u\) 自身が担当する道路の総数 \(m_u\) は、 - \(dp_0[u]\) の場合(親との道路は親が担当):\(m_u = c\) - \(dp_1[u]\) の場合(親との道路は \(u\) が担当):\(m_u = c + 1\)

となります。それぞれのケースについて、集落 \(u\) 自身で発生する追加コスト \(W_u \times \max(0, m_u - C_u)\) を加算し、すべての \(c\) の中での最小値を取ることで \(dp_0[u]\) および \(dp_1[u]\) を更新できます。


アルゴリズム

  1. グラフの構築: 与えられた親の関係から、各頂点の子リストを構築します。
  2. DPの実行:
    • 葉(子を持たない頂点)から順に根に向かって処理を行います。制約より \(P_i < i\) であるため、頂点番号を \(N\) から \(1\) まで逆順にループすることで、トポロジカル順(ボトムアップ)に処理できます。
    • 各頂点 \(u\) について:
      1. 子集落 \(v\) ごとに差分 \(D[v] = dp_0[v] - dp_1[v]\) を計算し、昇順にソートします。
      2. \(c = 0\) から \(c = k\)(子の総数)までループを回し、累積和を利用して「子を \(c\) 個選んだときの最小コスト」を求めます。
      3. \(dp_0[u]\)\(dp_1[u]\) のそれぞれのケースについて、自身が担当する道路数 \(m_u\) に応じた追加コストを加算し、最小値を記録します。
  3. 答えの出力: 根(頂点1)には親がいないため、根の上の道路を「親が担当する(=存在しない)」ものとした \(dp_0[1]\) が、追加コストの最小値となります。これに基本工事コスト \(N - 1\) を足した値を出力します。

計算量

時間計算量: \(O(N \log N)\)

各頂点 \(u\) において、子の個数を \(k_u\) とすると、ソートに \(O(k_u \log k_u)\) の時間がかかります。 すべての頂点に対する \(k_u\) の総和は、木の辺の数である \(N - 1\) です。 したがって、全体のソートにかかる時間は \(\sum_{u=1}^{N} O(k_u \log k_u) \leq O(N \log N)\) となり、十分に高速です。

空間計算量: \(O(N)\)

木構造を表現する隣接リスト、およびDPテーブル(dp0, dp1)の保持に必要なメモリは \(O(N)\) です。


実装のポイント

  • 再帰を避けたボトムアップ処理: Pythonなどの言語では、再帰関数を用いた深さ優先探索(DFS)を行うと、再帰の最大深度制限に達したり、オーバーヘッドが大きくなったりすることがあります。本問題では親の番号が自身の番号より小さい(\(P_i < i\))という性質があるため、ループを N から 1 まで逆順に回すだけで、再帰を使わずに安全かつ高速にボトムアップな木DPを実装できます。

  • 基本工事コストの加算: 問題文にある「基本工事コスト \(1 \times (N-1)\)」を最後に足し忘れないように注意してください。

    ソースコード

import sys


def solve():
    input = sys.stdin.read
    data = input().split()
    if not data:
        return

    N = int(data[0])
    P = [0, 0] + [int(x) for x in data[1:N]]

    C = [0] * (N + 1)
    W = [0] * (N + 1)
    idx = N
    for i in range(1, N + 1):
        C[i] = int(data[idx])
        W[i] = int(data[idx + 1])
        idx += 2

    children = [[] for _ in range(N + 1)]
    for i in range(2, N + 1):
        children[P[i]].append(i)

    dp0 = [0] * (N + 1)
    dp1 = [0] * (N + 1)

    for u in range(N, 0, -1):
        ch = children[u]
        k = len(ch)
        if k == 0:
            dp0[u] = 0
            diff = 1 - C[u]
            dp1[u] = W[u] * (diff if diff > 0 else 0)
            continue

        D = [dp0[v] - dp1[v] for v in ch]
        D.sort()

        sum_dp1 = 0
        for v in ch:
            sum_dp1 += dp1[v]

        curr_S = sum_dp1
        diff_C0 = -C[u]
        diff_C1 = 1 - C[u]
        min_val0 = curr_S + W[u] * (diff_C0 if diff_C0 > 0 else 0)
        min_val1 = curr_S + W[u] * (diff_C1 if diff_C1 > 0 else 0)

        for c in range(1, k + 1):
            curr_S += D[c - 1]
            diff0 = c - C[u]
            diff1 = c + 1 - C[u]
            val0 = curr_S + W[u] * (diff0 if diff0 > 0 else 0)
            val1 = curr_S + W[u] * (diff1 if diff1 > 0 else 0)
            if val0 < min_val0:
                min_val0 = val0
            if val1 < min_val1:
                min_val1 = val1

        dp0[u] = min_val0
        dp1[u] = min_val1

    ans = (N - 1) + dp0[1]
    print(ans)


if __name__ == "__main__":
    solve()

この解説は gemini-3.5-flash-thinking によって生成されました。

posted:
last update: