問題
$N$ 頂点の木があり、頂点 $i$ には(負にもなる)重み $w_i$ がある。ちょうど $K$ 個の頂点からなる連結な部分集合(誘導部分木)を選び、重みの総和を最大にしたい。その最大値を求めよ。
入力形式
N K
w_1 w_2 ... w_N
u_1 v_1
...
u_{N-1} v_{N-1}
制約
| # | 条件 |
|---|---|
| 1 | $1\le K\le N\le 5000$ |
| 2 | $|w_i|\le10^9$ |
| 3 | 与えられるグラフは木 |
| 4 | 入力はすべて整数 |
入出力例
入力例1
5 3
1 -2 3 4 -1
1 2
1 3
2 4
2 5出力例1
3($\{1,2,4\}$ で $1-2+4=3$。$\{1,3,2\}$ は 2)
入力例2
4 4
-1 -1 -1 -1
1 2
2 3
3 4出力例2
-4入力例3
4 3
1 1 1 1
1 2
1 3
1 4出力例3
3背景と出典
「木 DP の併合で、部分木のサイズを $K$ で打ち切ると全体で $O(NK)$ になる」は競プロで広く知られた性質で、テキストでは『木ナップサック』と呼ばれる、と私は記憶している(本セッションでは原典を開いていない)。
この問題の正しさの根拠は、全部分集合を列挙する全探索との照合(検証ログ)。$O(NK)$ の上界の証明は本文(Step 4)に議論を書いたが、形式的な証明ではない。実測では $K$ が小さいと速いこと($N=5000$、$K=10$ で 0.05 s、$K=5000$ で 0.84 s)だけを確かめた。
ヒント(クリックで段階的に開示)
ヒント1(方向性)
頂点 $v$ を根とする部分木の中で、$v$ を含む連結集合に限ると、子の側は「その子を含む連結集合」を $v$ にくっつけるか、くっつけないかの 2 択になる。この形なら、子の表を 1 つずつ併合していける。
ヒント2(アプローチ)
$f[v][j]$ = $v$ を含み $v$ の部分木に収まる連結な $j$ 頂点集合の重み最大。初期値 $f[v][1]=w_v$。子 $c$ を併合するとき、新しい $f[v][a+b]=\max(f[v][a]+f[c][b])$($b\ge1$)。子を使わない場合は $f[v][a]$ が残る。
ヒント3(誘導)
答えは根ではなく、全頂点 $v$ の $f[v][K]$ の最大。連結集合の最も根に近い頂点を $v$ とみなすと、各連結集合はちょうど 1 つの $v$ の $f[v][K]$ に数えられる。併合のとき $a+b\le K$ の範囲だけ見ると、1 回の併合は $\min(\text{size}_v,K)\cdot\min(\text{size}_c,K)$ になる。
---
模範解答 (Python)
模範解答を表示する(まず自力で解いてから)
import sys
NEG = -(1 << 60)
def main():
data = sys.stdin.buffer.read().split()
n, K = int(data[0]), int(data[1])
w = [int(x) for x in data[2:2 + n]]
adj = [[] for _ in range(n)]
p = 2 + n
for _ in range(n - 1):
u = int(data[p]) - 1; v = int(data[p + 1]) - 1; p += 2
adj[u].append(v); adj[v].append(u)
# 根 0 から順序と親を求める(再帰を使わない)
parent = [-1] * n; order = [0]
for v in order:
for c in adj[v]:
if c != parent[v]:
parent[c] = v; order.append(c)
# f[v][j] = v を含み部分木内に収まる連結な j 頂点集合の重み最大(j=1..min(size,K))。添字 0 は未使用
f = [None] * n
best = NEG
for v in reversed(order):
cur = [NEG, w[v]] # j=1 は v だけ
for c in adj[v]:
if c == parent[v]: continue
fc = f[c]; f[c] = None # 使い終わった子の表は捨ててメモリを抑える
# 子を使わない(そのまま)か、子の側から b>=1 頂点を足す
new = cur + [NEG] * (min(len(cur) - 1 + len(fc) - 1, K) - (len(cur) - 1))
for a in range(1, len(cur)):
ca = cur[a]
if ca == NEG: continue
lim = min(len(fc) - 1, K - a)
for b in range(1, lim + 1):
t = ca + fc[b]
if t > new[a + b]: new[a + b] = t
cur = new
f[v] = cur
if len(cur) > K and cur[K] > best: best = cur[K]
print(best)
main()
計算量: $O(NK)$(Step 4 の議論による。形式的な証明ではない)、空間は子の表を使い終わった時点で捨てるので同時に持つ表の合計は $O(N)$ 程度。実測(Python 3.14.7、M1 Mac、$|w|\le10^9$)で $N=5000$: $K=10$ のランダム木 0.05 s、$K=100$ で 0.08 s、$K=5000$ で 0.84 s、$K=2500$ のパスで 0.75 s、$K=2500$ のスターで 1.76 s。
---
Step-by-Step 解説
頂点 1 を根にして、親と処理順(order)を再帰なしで作る。パスの形の木($N=5000$)で再帰すると深さが足りなくなるので、order を後ろから処理する形にする。
$f[v][j]$ は「$v$ を含み、部分木に収まる」連結集合に限る。$v$ を含まない連結集合は、その集合の最上位の頂点 $u$ の $f[u][\cdot]$ に入っている。したがって答えは全頂点の $f[v][K]$ の最大(根だけを見ると $v$ が根でない最適解を落とす)。
子 $c$ を 1 つずつ $v$ に併合する。子を使わないなら $f[v][a]$ のまま、使うなら $c$ を含む連結集合 $b\ge1$ 個を足して $a+b$。$a+b\le K$ の範囲だけ作る。new = cur + [NEG] * ... で配列を伸ばしてから更新する。
サイズを $K$ で打ち切った併合のコストは $\min(s_1,K)\cdot\min(s_2,K)$。これを全併合で足すと、各頂点の組が最大 1 回ずつ、しかも一方が $K$ 以上の部分木に属する組の数は $O(NK)$ という議論になる。形式的な証明は教科書にあるはずだが、ここでは書かない(上の「背景と出典」)。
$K=1$ は $\max w_v$($f[v][1]=w_v$ が答え)。$K=N$ は全頂点の和。どちらも照合した 400 件の中に入っている($K$ は $1\le K\le n$ の全範囲から無作為に選んだ)。
---
迷った他案と採らなかった理由
| 案 | 計算量 | 採否 | 理由 |
|---|---|---|---|
| $f[v][j]$($v$ を含む連結集合)で併合を $K$ で打ち切る(本解) | $O(NK)$ | 採用 | 全探索と 400 件全件一致 |
| 全部分集合を列挙して連結性を確かめる | $O(2^N N)$ | 検算用のみ | $N\le11$ でしか動かない |
| 答えを $f[\text{根}][K]$ だけにする | $O(NK)$ | 不採用 | 400 件中 113 件で答えが小さい(根を含まない最適解を落とす) |
| 打ち切りなしで全サイズを持つ | $O(N^2)$ | 不採用 | $K$ が小さいとき無駄。実測は $N=5000,K=10$ で 0.05 s、$K=5000$ で 0.84 s |
| 最大重み連結部分木($K$ 指定なし)の貪欲 | $O(N)$ | 範囲外 | 本問は $K$ を指定する |
前提が変わったら:
- $N\ge10^5$ かつ $K$ も大きいと $O(NK)$ が届かない。重みが凸な構造を持つ特殊な場合(例えば Day 162 Q3 の凹畳み込み)だけ速くできる
- 重みが辺にある(辺の和)場合は、$K$ 頂点なら $K-1$ 辺。状態を辺側に直せば同じ形
- 連結でなくてもよいなら、重みの大きい順に $K$ 個選ぶだけ(木構造が不要)
---
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| 答えを根の $f[\text{根}][K]$ にする | 根を含まない最適解を落とす | 全頂点の $f[v][K]$ の最大(113/400 誤答) |
| 併合を全サイズで回す | $K$ 以上は不要なのを忘れる | lim = min(len(fc)-1, K-a) |
| 子の表を保持し続ける | 使い終わった子の表を None にしない | f[c] = None |
| 再帰 DFS で書く | 木がパスのとき深さが $N$ | order を作って逆順に処理 |
次のステップ
発展1: 頂点をちょうど $K$ 個でなく高々 $K$ 個にする($\max_j f[v][j]$)。重みが負のとき違いが出る
- 発展2: 根を固定した版(頂点 1 を必ず含む)にして、全方位 DP で全根の答えを出す
- 次回予告: 焦点帯外(M2・M3)の問題
検証ログ(本セッションで実行した数値)
- 全探索との照合(実測): $n\le11$、$|w|\le9$ の 400 件(ランダム木・パス・スターを 3 種類均等に)で、全部分集合を列挙して連結性を判定する全探索と全件一致
- 誤った方法(実測): 答えを $f[\text{根}][K]$ だけにする方法は 400 件中 113 件で誤答
- 入力例(実測): 例 1 は 3、例 2 は -4、例 3 は 3(出力例どおり)
- 実行時間(実測): $N=5000$ で、$K=10$ のランダム木 0.05 s、$K=100$ 0.08 s、$K=5000$ 0.84 s、$K=2500$ のパス 0.75 s、$K=2500$ のスター 1.76 s。$N=2000,K=1000$ はランダム木 0.14 s
- 扱っていない範囲: $O(NK)$ の形式的な証明、$N\ge10^5$、連結でない場合、重みが辺にある場合