問題概要

$n (\leq 200)$ 頂点の木があり，$i$ 番目の頂点の重みは$a_i (\leq 10^5)$ である．

この木の頂点の集合で，含まれるすべての頂点対間の距離が k を超えるという条件の下で，頂点集合中の重みの和を最大化せよ．

解法概要

空間計算量が $O(n^3)$ となる解のほかに，さらに空間計算量を削減した $O(n^2)$ 解も掲載する．

$O(n^3)$ 解

部分木の中での解を合成しながら葉から根に向かって上がっていく木 DP を考える．

$\mathit{dp}_{v, d}$: 頂点 v を根 (深さ0) とし，重みに和に使った頂点の中で深さが最小のものの深さが d であるような場合の重みの和の最大値

f:id:noimin:20191023221559p:plain

この DP の更新の仕方を考えてみる．例えば，上の図で k = 2 として $\mathit{dp}_{v, 1}$ を考えてみる．白抜きのノードを使おうとすると，頂点 v の子孫うち白抜きのノードの方向以外の子孫は， v からみた深さが 2 以上必要になる．同様に考えていくと，使う子孫の中で最も深さの浅い子孫を1つ決めうちし，他の方向の子孫に必要な深さは決め打った子孫の深さに応じて定めるという方針になる．式で表すと，

$\mathit{dp}_{v, d} = \max_{c \in \mathit{child}(d)} (\mathit{dp}_{c, d-1} + \sum_{c_2 \in \mathit{child}(d), c_2 \neq c} \max_{\max (d-1, k-d) \leq d_c} \mathit{dp}_{c_2, d_c} )$

である．ただし d = 0 の場合は

$\mathit{dp}_{v, 0} = \sum_{c \in \mathit{child}(d)} \max_{k \leq d_c} \mathit{dp}_{c, d_c}$

となる．

この解法は

dfsで全ノードを見る
各ノードについて，ありうる$d$ の値をすべて見る
深さが最小の子を決め打つ
深さが最小の子以外の子について和を求める

といった処理があることから，空間計算量は一見 $O(n^4)$ のようにも見える．しかし，このうち「dfs で全ノードを見る」と「深さが最小の子を決め打つ」のみに着目すると，直接親子関係を持っているようなノードのペア$(\mathit{parent}, \mathit{child})$ の数は $(n-1)$ である (木構造なので) ．そのため空間計算量も $O(n^3)$ で上から抑えることができる．

$O(n^2)$ 解

$O(n^3)$ 解の $\max_{\max (d-1, k-d) \leq d_c} \mathit{dp}_{c_2, d_c}$ および $\max_{k \leq d_c} \mathit{dp}_{c, d_c}$ は，return の前に $d$ が大きい順から $\mathit{dp}_{v, d} $ を見ていくことで事前に計算しておくことが可能である．この事前計算を行うことで空間計算量を $O(n^2)$ で抑えることができる．

ちなみに私の提出では 389 ms ($O(n^3)$ 解) → 31ms ($O(n^2)$ 解) のように実行時間を縮めることができた．

ソースコード

    #include <iostream>
    #include <vector>
    #include <algorithm>
     
    using namespace std;
     
    using ll =  long long;
     
    constexpr int N_MAX = 210;
     
    int n, k;
    vector<int> graph[N_MAX];
    vector<ll> w(N_MAX, 0);
    vector<vector<ll>> dp(N_MAX, vector<ll>(N_MAX, 0));
     
    void dfs(int node, int par) {
        for(int child: graph[node]) {
            if(child == par) continue;
            dfs(child, node);
        }
     
        for(int d=0;d<N_MAX;++d) {
            if(d == 0) {
                dp[node][d] = w[node];
                for(int child: graph[node]) {
                    if(child == par) continue;
                    dp[node][0] += dp[child][k];
                    // O(n^3) 解
                    // dp[node][0] += *max_element(dp[child].begin()+k, dp[child].end());
                }
            } else {
                for(int child: graph[node]) {
                    if(child == par) continue;
                    ll tmp = dp[child][d-1];
                    for(int child2: graph[node]) {
                        if(child2 == par || child2 == child) continue;
                        tmp += dp[child2][max(d-1, k-d)];
                        // O(n^3) 解
                        // tmp += *max_element(dp[child2].begin()+max(d-1, k-d), dp[child2].end());
                    }
                    dp[node][d] = max(dp[node][d], tmp);
                }
            }
        }
     
        // O(n^3) 解では不要
        for(int d=N_MAX-2;d>=0;--d) {
            dp[node][d] = max(dp[node][d], dp[node][d+1]);
        }
    }
     
    int main() {
        ios::sync_with_stdio(false); cin.tie(nullptr);
     
        int a, b;
        cin >> n >> k;
        for(int i=0;i<n;++i){
            cin >> w[i];
        }
        for(int i=0;i<n-1;++i){
            cin >> a >> b;
            --a; --b;
            graph[a].push_back(b);
            graph[b].push_back(a);
        }
     
        dfs(0, -1);
        cout << *max_element(dp[0].begin(), dp[0].end()) << endl;
    }

所感

Div. 3 初参加でした． Speedrun 感覚で参加できて楽しい．

と言いつつ，この問題は解説 AC だけど．

それにしても最近こどふぉの問題の解説ばっかり書いてるな．

NoiminのNoise

Codeforces Round #595 Div. 3 F. Maximum Weight Subset

問題概要

解法概要

$O(n^3)$ 解

$O(n^2)$ 解

ソースコード

所感