問題
$N$ 頂点の根付き木(根は頂点1)があり、頂点 $i$ には価値 $v_i$(負の値もありうる)が定められている。頂点の部分集合 $S$ を選ぶ。ただし根は必ず $S$ に含め、頂点 $u\in S$($u\ne1$)ならその親も $S$ に含まれていなければならず、$|S|\le K$ とする。$\sum_{i\in S} v_i$ の最大値を求めよ。
入力形式
N K
v_1 v_2 ... v_N
p_2 p_3 ... p_N
制約
$1 \le K \le N \le 2000$
$|v_i| \le 10^9$
$1 \le p_i < i$
入出力例
入力例1
5 3
3 5 1 8 2
1 1 2 2
出力例1
16
頂点2,3の親は1、頂点4,5の親は2。{1,2,4}を選ぶと 3+5+8=16 が最大。単純に価値の大きい順に{4,2,5}を選ぶことはできない(4を選ぶには2、2を選ぶには1が必須)。
概念図: 依存関係付き部分木選択
ヒント(段階的開示)
ヒント1: 方向性
「価値が大きい頂点から貪欲に選ぶ」は失敗する。頂点4(価値8)を選ぶには先祖である頂点2と頂点1も同時に選ぶ必要があり、単純な貪欲では依存関係を無視してしまう。木の形を反映したDPが必要になる。
ヒント2: アプローチ
各頂点vについて「vを根とする部分木からvを必ず含めてj個選んだときの価値の最大値」をdp[v][j]として定義し、葉から根に向かって子の情報をマージしていく(子の部分木ごとに「何個選ぶか」を独立に決めてよい選択問題)。配列を「その時点での部分木サイズとKの小さい方」に切り詰めながらマージすると、全体の計算量は部分木サイズの積の総和の議論によりO(NK)に収まる。
ヒント3: 誘導(コード骨格)
dp = [None] * N # dp[v]は長さmin(size(v),K)+1のリスト
for v in 帰りがけ順:
cur = [NEG, value[v]] # vだけ選んだ状態
size_so_far = 1
for c in children[v]:
# curとdp[c]をナップサック的にマージ
new_cap = min(size_so_far + (len(dp[c]) - 1), K)
...
size_so_far = new_cap
dp[v] = cur
模範解答 (Python)
import sys
def solve():
data = sys.stdin.buffer.read().split()
idx = 0
N = int(data[idx]); idx += 1
K = int(data[idx]); idx += 1
value = [0] * N
for i in range(N):
value[i] = int(data[idx]); idx += 1
parent = [0] * N
for i in range(1, N):
p = int(data[idx]); idx += 1
parent[i] = p - 1
children = [[] for _ in range(N)]
for v in range(1, N):
children[parent[v]].append(v)
NEG = float('-inf')
post = []
stack = [(0, False)]
while stack:
node, processed = stack.pop()
if processed:
post.append(node)
else:
stack.append((node, True))
for c in children[node]:
stack.append((c, False))
dp = [None] * N
for v in post:
cur = [NEG, value[v]]
size_so_far = 1
for c in children[v]:
child_dp = dp[c]
new_size_cap = min(size_so_far + (len(child_dp) - 1), K)
new_cur = [NEG] * (new_size_cap + 1)
for j1 in range(0, min(size_so_far, K) + 1):
if cur[j1] == NEG:
continue
if j1 <= new_size_cap and new_cur[j1] < cur[j1]:
new_cur[j1] = cur[j1]
for j2 in range(1, len(child_dp)):
if child_dp[j2] == NEG:
continue
tot = j1 + j2
if tot > new_size_cap:
continue
val = cur[j1] + child_dp[j2]
if val > new_cur[tot]:
new_cur[tot] = val
cur = new_cur
size_so_far = new_size_cap
dp[v] = cur
root_dp = dp[0]
print(max(x for x in root_dp if x != NEG))
solve()
計算量: 部分木サイズの積の総和の議論により全体O(NK)(実質O(N min(N,K)))。N≤8のランダムな木を200通り生成し、全列挙による厳密解と本DPの出力を突き合わせて全一致を確認済み。
Step-by-Step 解説
1DPの定義
dp[v][j]を「vを根とする部分木内からvを必ず含めてちょうどj個選んだ最大価値」とする。答えはmax_j dp[1][j](j≤K)。
dp[v][j]を「vを根とする部分木内からvを必ず含めてちょうどj個選んだ最大価値」とする。答えはmax_j dp[1][j](j≤K)。
2子1つずつマージするナップサック
子ごとに「使わない」か「j2個使う」かを選べるグループナップサックの形。cur[j1]とchild_dp[j2]を全ペア試してnew_cur[j1+j2]を更新する。
子ごとに「使わない」か「j2個使う」かを選べるグループナップサックの形。cur[j1]とchild_dp[j2]を全ペア試してnew_cur[j1+j2]を更新する。
3なぜ全体O(NK)に収まるか
配列を毎回「部分木サイズとKの小さい方」に切り詰めていると、マージのコストは部分木サイズの積の総和の議論によりO(N・min(N,K))に抑えられる。
配列を毎回「部分木サイズとKの小さい方」に切り詰めていると、マージのコストは部分木サイズの積の総和の議論によりO(N・min(N,K))に抑えられる。
4「使わない」選択肢を忘れない
子の部分木を丸ごと使わない選択肢を明示的にコピーしないと、常に全ての子を最低1個使わなければならない誤った制約になる。
子の部分木を丸ごと使わない選択肢を明示的にコピーしないと、常に全ての子を最低1個使わなければならない誤った制約になる。
5動作確認
N≤8のランダムな木200通りで全列挙の厳密解と突き合わせ全一致を確認済み。
N≤8のランダムな木200通りで全列挙の厳密解と突き合わせ全一致を確認済み。
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| 子の部分木を「使わない」選択肢を実装し忘れる | マージループでj2>=1の分岐しか書かない | new_cur[j1]=max(new_cur[j1],cur[j1])を先に必ず入れる |
| 配列サイズを毎回K+1固定で確保する | 部分木サイズで切り詰める最適化を怠る | min(部分木サイズ,K)+1に切り詰め計算量をO(NK)に抑える |
| 価値が負の頂点を除外して考えてしまう | 「正のものだけ選べばよい」と誤解する | 祖先である以上子孫を選ぶために選ばざるを得ない場合がある |
| dp[v][0]を有効な値として扱ってしまう | 「vを選ばない」状態と定義を混同する | 本解法の定義ではdp[v][0]は使わない(j=1が最小) |
次のステップ
- 発展: 「K個以内」ではなく「ちょうどK個」を選ぶ場合や、辺にもコストがある場合(頂点数でなく総コスト制約)に拡張してみる。
- 次回予告: 2次元BIT 矩形加算・矩形和クエリ(4-BITトリックの2次元拡張)