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)\) で答えることも可能です。
投稿日時:
最終更新:
