PythonでABC259F

Upsolveしたので、理解を強固にするために解説記事を書きました。

問題へのリンク

https://atcoder.jp/contests/abc259/tasks/abc259_f

提出コード

https://atcoder.jp/contests/abc259/submissions/33278369



## 問題のイメージ

$N$個の頂点と$N-1$本の辺からなる木があります。0本以上の辺を選んだ時重みの最大値はいくつでしょうか。ただし、各頂点$i$において選んだ辺の次数が$d_i$を超えてはいけません。

辺を蛍光灯に見立てて、それぞれの辺についてON/OFFを決めていくことを考えます。$d_i$のルールを守りながら、ON辺の重みの総和が最大になるようにしていきます。

以下では、OFF辺を黒色の点線、ON辺を赤色の実線で示します。
![graph2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/253018/add26d16-4237-2ade-2d20-a9a157e14529.png)


サンプル１をもとに、この問題でやろうとしていることをイメージしてみます。ここでは、頂点を0-indexに変えてあります。<>内の数字は$d_i$です。

最初、すべての辺はOFFです。

![graph.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/253018/16e11886-aac5-e866-bdbe-97a08ec6f9d4.png)

このときの正答は↓になります。

![graph.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/253018/08d253b4-a3d2-801c-bcf3-a3b53d202528.png)

ON辺は$0$本でもよいし、$N-1$本すべてでもよいです。ON辺のみで木をなしている必要はなく、途中で途切れていてもよいです。

それでは、最適解を導き出すアルゴリズムはどのようになるでしょうか？

## 考察

#### 部分木から考える

全体の最適解を一発で決めるのが難しいときは、小さなパーツに分けて考えてから徐々に大きくしていくのが定石です。ということで、部分木について思いを馳せます。

下図でいうと青→緑→赤のように、木の部分木、その部分木のそのまた部分木と、どんどん細かくしていくと、最終的に葉のところでは頂点1つだけの部分木になります。この最小単位の部分木における重み総和の最大値は必ず0です（辺を持たないのため）。これで最も小さなパーツの値が定ままりました。これを足掛かりにして子から親へとさかのぼっていきましょう。

![graph03.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/253018/ec026a40-7488-f996-f1ce-d6434d31ebba.png)


#### 2種類の状態

「子頂点の部分木」から「親頂点の部分木」への遷移の仕方を考えます。

例として、以下のような木の重み総和の最大値を求めてみます。

![graph04.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/253018/63cfe789-296f-c652-aa86-966256e8885e.png)

さきほどの方針通り、子から親への方向へ「重み総和の最大値」を決めていきます。
赤色部分の重み総和の最大値はそれぞれ0です。
緑色エリア（頂点1以下の部分木）の重み総和の最大値を考えます、$d_i$の許すかぎり貪欲に辺をOFFからONに変えていくと、3本ある辺すべてをONにできて、最大値は20+30+40=90です。
次に、青色エリア（木全体）を見てみます。頂点1にはすでにON辺が3本つながっているので0と1をつなぐ辺を選ぶことができません。しかし、あきらかに重み100である辺をONにしたほうが重み総和の最大値は大きくなります。実際のこの木の最適解は30+40+100=170です。

このケースが示唆することは、「頂点$v$以下の部分木における、$d_v$本の辺をすべて使いきった時の重み総和の最大値」だけを保持していては「親頂点の部分木における重み総和の最大値」を求めるには不十分だということです。辺をONにできる権利を1本残しておいて、頂点$v$とその親を結ぶ辺をONにできるようにしておいたときのほうが良い場合もあります。

![graph05.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/253018/77e76ebc-1a18-4d7e-1662-181a8c95f332.png)

左：親への辺を残しておかない場合、右：親への辺をONにする権利を残しておき、かつ、その権利を使って親への辺をONにした場合


#### 木DP

ここまでの考察から、各頂点以下の部分木において2つの状態を保持しながら、子から親へと遷移させていきます。具体的には、以下のような2つの状態を保持してDPをしていきます。

```math

\begin{align}
dp[v][0] &= 頂点v以下の部分木について、vに接続する辺をd_v本以下ONにしたときの重み総和の最大値\\

dp[v][1] &= 頂点v以下の部分木について、vに接続する辺をd_v-1本以下ONにしたときの重み総和の最大値
\end{align}
```

$dp[v][0]$が親への辺を残しておかない状態、$dp[v][1]$が1本節約して親への辺をONにする余裕を残しておいた状態に対応します。この2つの状態の大小関係は常に$dp[v][0] \geq dp[v][1]$になります。

特殊な例として、$d_v=0$のとき、辺を1本残しておくということは不可能ですから $dp[v][1]=-\infty$ としておきます。

#### 辺の選び方

最後に考えるのは、複数の辺が存在する場合に、どのように優先度を付けて辺をONにすればよいのかということです。以下のようなケースを考えます。

```text
15
2 0 2 2 2 2 1 1 1 1 1 1 1 1 1
1 2 50
1 3 -100
1 4 50
1 5 50
1 6 50
1 7 50
3 8 60
3 9 70
4 10 70
4 11 60
5 12 30
5 13 40
6 14 10
6 15 20
```


![graph06.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/253018/bee43a49-0f40-3fba-e425-446fcd3d111f.png)

頂点0を根として、この木は深さ0,深さ1,深さ2の階層に分けられます。深さ1の頂点1,2,3,4,5,6のdp[v]の計算まではすでに終わっているとして、頂点0とその子頂点$u_i=1,2,3,4,5,6$をつなぐ重み$w_i$の辺を選んでONにしていくプロセスを詳しく見てみます。

はじめ、辺6本はすべてOFFです。
辺(0,1)はONにできません。頂点1が$d_v=0$であるため。
辺(0,2)の重みは負なので、この辺をONにすると損します。
辺(0,3)をONにすると、かわりに辺(3,9)をOFFにしなければいけません。60を捨てて50を拾うため10だけ損します。
辺(0,4)をONにすると、かわりに辺(4,11)をOFFにしなければいけません。30を捨てて50を拾うため20得します。
辺(0,5)をONにすると、かわりに辺(4,11)をOFFにしなければいけません。10を捨てて50を拾うため40得します。
辺(0,6)をONにすると、単純に50得します。


これらの観察から、辺をOFFからONに変えたときは$dp[u_i][0]$を捨てて$dp[u_i][1]+w_i$を拾うことがわかります。
辺をONにしたときのお得度を

```math
\Delta_i=-dp[u_i][0]+(dp[u_i][1]+w_i)
```

と定義するとよさそうです。このように定義すると、辺をONにすることができないときやONにすると損するときは$\Delta_i<0$になります。

|$u_i$|1|2|3|4|5|6|
|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
|$dp[u_i][0]$                            |0        |130  |130 |70|30|0|
|$dp[u_i][1]$                            |$-\infty$|70 | 70 |40|20|0|
|$w_i$                                   |      0  |-100|50  |50|50|50|
|$\Delta_i=-dp[u_i][0]+(dp[u_i][1]+w_i)$|$-\infty$ | -160|-10|20|40|50|

dp[v][0]を得るには、お得度$\Delta_i$を降順ソートして、その順番で、$d_v$本以下かつ$\Delta_i>0$であるかぎり辺をONにしていくのがよいです。今のケースだとお得度は[6,5,4,3,2,1]の順番になり、$d_0=2$であることから辺(0,6),(0,5)をONにするのが最適です。もし$d_0=4$であるときは、辺(0,6),(0,5),(0,4)をONにするのが最適です（辺(0,3)はお得度が負のためONにしません）。
同様に、dp[v][1]を得るには、お得度$\Delta_i$を降順ソートして、その順番で、   $d_v$-1 本以下 かつ$\Delta_i>0$であるかぎり辺をONにしていくのがよいです。



最終的にdp[0][0]が答えとなります。今ケースの答えは450です。

![graph06.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/253018/40214eee-1a70-1540-e247-000518e79d7f.png)

## 実装

アルゴリズムを実現するには、まず、葉から根へ向かう順番で頂点にアクセスする必要があります。そうしないと、ある頂点にアクセスしたときにその子頂点のdpの値が未定である事態が発生します。2つのやり方があります。

1. 幅優先探索を行って「行きがけ順」をメモし、その逆順に頂点へアクセスする（＝深さが大きい頂点から順に決めていく）。
1. 深さ優先探索を用いて木の末端から決めていく。

今回は1を使います。

実装ソースコードは下記です。わかりやすさ重視のため記述が冗長になっていることがあります。

```python
from collections import deque
N = int(input())
d = list(map(int, input().split()))
graph = [[] for _ in range(N)]

for _ in range(N-1):
    a, b, w = map(int, input().split())
    a, b = a-1, b-1
    graph[a].append((b, w))
    graph[b].append((a, w))

#頂点0を根とする
start = 0

#訪問したかどうか
is_visited = [False]*N
is_visited[start] = True

#訪問した順
order_visited = []
order_visited.append(start)

#親頂点をメモしておく
parent_node = [-1]*N

#幅優先探索で行きがけ順を記録する
q = deque()
q.append(start)

while q:
    v = q.popleft()
    for u, _ in graph[v]:
        if not is_visited[u]:
            is_visited[u] = True
            order_visited.append(u)
            parent_node[u] = v
            q.append(u)

#木DP
dp = [[0, 0] for _ in range(N)]
INF = 10**18

for v in reversed(order_visited):
    
    #子頂点それぞれについてお得度Deltaを計算する
    chile_info = []

    for u, w in graph[v]:
        if u == parent_node[v]:
            continue

        chile_info.append([u, w, dp[u][1]+w-dp[u][0]])

    #お得度で降順ソート
    chile_info.sort(reverse=True, key=lambda x: x[2])

    #dp[v][0]を計算（親への辺を残しておかない場合）
    for i, (u, w, delta) in enumerate(chile_info):
        if i+1 <= d[v] and delta > 0:
            dp[v][0] += dp[u][1]+w
        else:
            dp[v][0] += dp[u][0]

    #dp[v][1]を計算（親への辺を残しておく場合）
    for i, (u, w, delta) in enumerate(chile_info):
        if i+1 < d[v] and delta > 0:
            dp[v][1] += dp[u][1]+w
        else:
            dp[v][1] += dp[u][0]

    if d[v] == 0:
        dp[v][1] = -INF

print(dp[0][0])
```

## 参考文献

https://twitter.com/kyopro_friends/status/1545767640396079105?s=20&t=biLSLje2hPjDEN0F-0uGIg

https://www.youtube.com/watch?v=8pv0U8-Oe-Y&t=6832s