概要

木構造の通信ネットワークで、青木君が高々 \(1\) 本の辺を切った後、高橋君が距離 \(D\) 以内で盗めるデータ量の最大値を最小化する問題です。

各辺を切った場合について、「切られる前に盗めた量」から「切断によって届かなくなる反対側の量」を引く形で高速に計算します。

考察

まず、辺を切らない場合、各頂点 \(x\) から距離 \(D\) 以内にある頂点のデータ量の合計を

[ S_x ]

とします。

辺を切らない場合、高橋君は最も大きい \(S_x\) を選ぶので、盗める量は

[ \max_x S_x ]

です。

---

素朴に考えると、各辺を切るたびに、各頂点を侵入先として距離 \(D\) 以内の頂点を調べることになります。

辺の候補は \(O(N)\) 個、侵入先も \(O(N)\) 個、さらに到達可能な頂点を調べると \(O(N)\) かかるため、全体で

[ O(N^3) ]

となります。\(N \leq 3000\) では間に合いません。

---

重要な観察は、ある辺を切ったときに「失われる頂点」は、切った辺の反対側にある頂点だけだということです。

木を適当な頂点、ここでは頂点 \(1\) を根として根付き木にします。

ある辺を親 \(p\) と子 \(u\) を結ぶ辺とします。この辺を切ると、木は次の \(2\) つに分かれます。

• \(u\) を根とする部分木側
• それ以外の補集合側

ここで、侵入先 \(x\) が \(u\) 側にあるとします。

反対側の頂点 \(y\) に行くには、必ず辺 \((p,u)\) を通る必要があります。そのため、元の木での距離は

[

\operatorname{dist}(x,y)

\operatorname{dist}(x,u) + 1 + \operatorname{dist}(p,y) ]

です。

したがって、\(x\) から距離 \(D\) 以内にあった反対側の頂点は、

[ \operatorname{dist}(p,y) \leq D - 1 - \operatorname{dist}(x,u) ]

を満たす頂点です。

つまり、\(x\) が \(u\) からどれだけ離れているかだけで、切断によって失われる反対側のデータ量が決まります。

例えば \(D=5\) で、\(x\) が \(u\) から距離 \(2\) の場所にあるなら、反対側で届かなくなるのは

[ \operatorname{dist}(p,y) \leq 5 - 1 - 2 = 2 ]

を満たす頂点です。

---

よって、同じ側にある頂点を「切断辺の端点からの距離」でまとめればよいです。

\(u\) 側について、距離 \(t\) ごとに

• その距離にある頂点のうち最大の \(S_x\)
• 反対側で距離 \(D-1-t\) 以下にあるデータ量の合計

が分かれば、切断後に \(x\) から盗める最大値を計算できます。

補集合側についても同様です。

アルゴリズム

1. 木を根付き木にする

頂点 \(1\) を根として DFS し、各頂点について

• 親 parent
• DFS 順 order
• 部分木サイズ size

を求めます。

DFS 順を使うと、ある頂点 \(u\) の部分木は order 上で連続した区間になります。

これにより、辺 \((parent[u], u)\) を切ったとき、

• \(u\) の部分木側
• それ以外の側

を簡単に列挙できます。

---

2. 全頂点間距離と \(S_x\) を求める

各頂点 \(s\) から DFS して、すべての頂点への距離を求めます。

その距離が \(D\) 以下の頂点のデータ量を合計して、

[ Ss = \sum{\operatorname{dist}(s,v) \leq D} V_v ]

を計算します。

辺を切らない場合の答え候補は

[ \max_s S_s ]

です。

---

3. 各辺を切った場合を調べる

辺 \((p,u)\) を切るとします。

\(u\) 側を in、補集合側を out と呼びます。

in 側について

頂点 \(x\) が in 側にあり、

[ t = \operatorname{dist}(u,x) ]

とします。

もし \(t < D\) なら、切断によって届かなくなる out 側の頂点は

[ \operatorname{dist}(p,y) \leq D - 1 - t ]

を満たす頂点です。

そのデータ量の合計を引けば、切断後に \(x\) から盗める量になります。

[ S_x - \text{out側で }p\text{ から距離 }D-1-t\text{ 以下のデータ量} ]

\(t \geq D\) の場合、反対側へ行くには最低でも \(D+1\) 距離が必要なので、もともと反対側の頂点は距離 \(D\) 以内に入っていません。そのため盗める量はそのまま \(S_x\) です。

out 側について

同様に、頂点 \(x\) が out 側にあり、

[ t = \operatorname{dist}(p,x) ]

なら、切断によって届かなくなる in 側の頂点は

[ \operatorname{dist}(u,y) \leq D - 1 - t ]

を満たす頂点です。

---

4. 距離ごとにまとめる

各辺について、以下を作ります。

• in_ex[t] : in 側で \(u\) から距離 \(t\) の頂点のデータ量合計
• out_ex[t] : out 側で \(p\) から距離 \(t\) の頂点のデータ量合計
• md[t] : in 側で \(u\) から距離 \(t\) の頂点のうち、最大の \(S_x\)
• mo[t] : out 側で \(p\) から距離 \(t\) の頂点のうち、最大の \(S_x\)

その後、in_ex, out_ex は累積和にします。

すると、

[ \text{out側で距離 }k\text{ 以下のデータ量} ]

を \(O(1)\) で取得できます。

---

5. その辺を切った場合の高橋君の最善値を求める

in 側で高橋君が得られる最大値は、

[ \max_t \left(md[t] - out_ex[D-1-t]\right) ]

です。

同様に out 側では、

[ \max_t \left(mo[t] - in_ex[D-1-t]\right) ]

です。

辺を切った後、高橋君はどちらの連結成分にも侵入できるので、その辺を切った場合の結果は

[ \max(\text{in側の最大値}, \text{out側の最大値}) ]

です。

青木君はこれを最小化したいので、全ての辺について最小値を取ります。

計算量

• 時間計算量: \(O(N^2)\)
• 空間計算量: \(O(N^2)\)

全頂点間距離の計算に \(O(N^2)\) かかります。

また、各辺について全頂点を一度ずつ見て情報を作るため、これも全体で \(O(N^2)\) です。

距離配列を全て保存するため、空間計算量は \(O(N^2)\) です。

実装のポイント

• データ量の合計は最大で \(3000 \times 10^9 = 3 \times 10^{12}\) になるため、C++ なら long long が必要です。Python では整数が任意精度なのでそのままで問題ありません。

• 木の深さが \(N\) になる可能性があるため、コードでは再帰 DFS ではなくスタックを使った DFS にしています。

• DFS 順 order と部分木サイズ size を使うことで、部分木を連続区間として扱えます。

• 距離 \(D\) 以上の頂点は、反対側の頂点を失うことがないため、別枠で最大の \(S_x\) だけを管理しています。

  ソースコード

import sys

def main():
    data = list(map(int, sys.stdin.buffer.read().split()))
    N = data[0]
    D = data[1]
    V = data[2:2 + N]

    adj = [[] for _ in range(N)]
    idx = 2 + N
    for _ in range(N - 1):
        a = data[idx] - 1
        b = data[idx + 1] - 1
        idx += 2
        adj[a].append(b)
        adj[b].append(a)

    parent = [-1] * N
    parent[0] = -2
    order = []
    st = [0]
    while st:
        v = st.pop()
        order.append(v)
        pv = parent[v]
        for to in adj[v]:
            if to != pv:
                parent[to] = v
                st.append(to)
    parent[0] = -1

    size = [1] * N
    for v in reversed(order):
        p = parent[v]
        if p >= 0:
            size[p] += size[v]

    dist = []
    S = [0] * N
    for s in range(N):
        row = [0] * N
        st = [s]
        pst = [-1]
        while st:
            v = st.pop()
            pr = pst.pop()
            nd = row[v] + 1
            for to in adj[v]:
                if to != pr:
                    row[to] = nd
                    st.append(to)
                    pst.append(v)

        total = 0
        for i in range(N):
            if row[i] <= D:
                total += V[i]
        S[s] = total
        dist.append(row)

    ans = max(S)

    NN = N
    DD = D
    last = DD - 1
    ordr = order
    SS = S
    VV = V
    distm = dist
    par = parent
    sz = size

    for oi in range(1, NN):
        u = ordr[oi]
        p = par[u]
        l = oi
        r = oi + sz[u]

        rowu = distm[u]
        rowp = distm[p]

        in_ex = [0] * DD
        out_ex = [0] * DD
        md = [-1] * DD
        mo = [-1] * DD
        md_deep = -1
        mo_deep = -1

        for ii in range(l, r):
            x = ordr[ii]
            t = rowu[x]
            sx = SS[x]
            if t < DD:
                in_ex[t] += VV[x]
                if sx > md[t]:
                    md[t] = sx
            elif sx > md_deep:
                md_deep = sx

        for ii in range(l):
            x = ordr[ii]
            t = rowp[x]
            sx = SS[x]
            if t < DD:
                out_ex[t] += VV[x]
                if sx > mo[t]:
                    mo[t] = sx
            elif sx > mo_deep:
                mo_deep = sx

        for ii in range(r, NN):
            x = ordr[ii]
            t = rowp[x]
            sx = SS[x]
            if t < DD:
                out_ex[t] += VV[x]
                if sx > mo[t]:
                    mo[t] = sx
            elif sx > mo_deep:
                mo_deep = sx

        for i in range(1, DD):
            in_ex[i] += in_ex[i - 1]
            out_ex[i] += out_ex[i - 1]

        sub = md_deep
        comp = mo_deep

        for t in range(DD):
            rev = last - t

            m = md[t]
            if m >= 0:
                val = m - out_ex[rev]
                if val > sub:
                    sub = val

            m = mo[t]
            if m >= 0:
                val = m - in_ex[rev]
                if val > comp:
                    comp = val

        cur = sub if sub >= comp else comp
        if cur < ans:
            ans = cur

    print(ans)

if __name__ == "__main__":
    main()

---

この解説は gpt-5.5-high によって生成されました。