K - 隣接禁止 / Not Adjacent 解説 by seekworser


与えられた木を頂点 \(0\) を根とする根付き木とみなします。

  • \(dp_{0, i, j}\):頂点 \(i\) の部分木のみ考えて、頂点 \(i\)選ばずに全部で \(j\) 個の頂点を問題文の条件を満たすように選ぶときの、選ばれた頂点に書かれた整数の総和の最大値
  • \(dp_{1, i, j}\):頂点 \(i\) の部分木のみ考えて、頂点 \(i\)必ず選び全部で \(j\) 個の頂点を問題文の条件を満たすように選ぶときの、選ばれた頂点に書かれた整数の総和の最大値

とする動的計画法を考えると、これは葉から順に計算可能です。この動的計画法の計算量は \(O(N^2 K)\) であり与えられた制約下で十分高速です。

#include <bits/stdc++.h>
using namespace std;
const int INF = 1 << 30;
int main() {
    int n,k; cin >> n >> k;
    vector g(n, vector<int>());
    for (int i=0; i<n-1; i++) {
        int u,v; cin >> u >> v; u--, v--;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    vector<int> a(n);
    for (int i=0; i<n; i++) cin >> a[i];
    auto dp = [&] (auto self, int u, int par) -> pair<vector<int>, vector<int>> {
        vector<int> dp0(k+1, -INF);
        vector<int> dp1(k+1, -INF);
        dp0[0] = 0;
        dp1[1] = a[u];
        for (int v : g[u]) {
            if (v == par) continue;
            auto [dp0_v, dp1_v] = self(self, v, u);
            vector<int> dp0_n(k+1, -INF);
            vector<int> dp1_n(k+1, -INF);
            for (int i=0; i<=k; i++) for (int j=0; j<=k; j++) {
                if (i+j <= k && dp0[i] != -INF && dp0_v[j] != -INF) dp0_n[i+j] = max(dp0_n[i+j], dp0[i] + dp0_v[j]);
                if (i+j <= k && dp0[i] != -INF && dp1_v[j] != -INF) dp0_n[i+j] = max(dp0_n[i+j], dp0[i] + dp1_v[j]);
                if (i+j <= k && dp1[i] != -INF && dp0_v[j] != -INF) dp1_n[i+j] = max(dp1_n[i+j], dp1[i] + dp0_v[j]);
            }
            swap(dp0, dp0_n);
            swap(dp1, dp1_n);
        }
        return {dp0, dp1};
    };
    auto [dp0, dp1] = dp(dp, 0, -1);
    int ans = max(dp0[k], dp1[k]);
    if (ans == -INF) ans = -1;
    cout << ans << '\n';
}

また、二乗の木DPと呼ばれる手法を用いて、この問題に \(O(N^2)\) で答えることも可能です。

投稿日時:
最終更新: