Ex - Many Illumination Plans Editorial by en_translator
We first consider the case where \(v=1\) and try to find \(F(1)\)).
Straightforward DP: \(\mathrm{O}(N X^2)\)
We first consider a naive tree DP (Dynamic Programming). We can come up with an algorithm costing \(\mathrm{O}(N X^2)\) time as follows.
- In the pseudocodes below,
c
,Dp[0]
andDp[1]
are used as temporary variables that mean:c
: currently inspected vertexDp[0], Dp[1]
: DP table of length \((X+1)\)
function dfs(c):
Dp[0] <- (DP table corresponding to empty set)
Dp[1] <- (DP table corresponding to empty set)
for d in {children of c} do
ep = dfs(d)
Dp[0] <- max_plus_convolution(Dp[0], ep[0])
Dp[1] <- max_plus_convolution(Dp[1], ep[1])
end for
for i = 0 to X - W[c] do
chmax(Dp[C[c]][i + W[c]], Dp[1 - C[c]][i] + B[c])
end for
return Dp
This algorithm costs \(\mathrm{O}(NX^2)\) time, because the max-convolution part costs \(\mathrm{O}(X^2)\) time. As \(X\) is very large, it is very hard to get AC (accepted).
Heavy-Light Decomposition Recursive DP, \(\mathrm{O}(N^{1.59} X)\)
We referred to an article and paper by tmaehara to write the following editorial.
This problem can be solved with a technique called Heavy-Light Decomposition Recursive DP (HLRecDP) in a total of \(\mathrm{O}(\mathrm{poly}(N) X)\) time.
The \(\mathrm{O}(NX^2)\) DP merges DP tables bottom-up for each subtree, but the problem is that merging costs \(\mathrm{O}(X^2)\).
On the other hand, consider computing the DP top-down to avoid merges. The DP table when the root is root
can be filled by calling the following recursive function with dfs(root, initial dp table)
. (This DP is quite irregular, requiring a smart thought. For some of you, it may help understanding to consider it as a DSU-on-tree algorithm (DSU=Disjoint set Union).)
- Note 1: In
dfs(c, dp)
,dp
is passed by reference. (As we are not modifyingdp
, it is sufficient to pass aconst
reference in C++.) - Note 2: Unlike bottom-up DP, note that the return value of the intermediate
dfs(c, dp)
is not “the answer to the subtree rooted atc
.”
function dfs(c, dp)
Dp[0] <- dp
Dp[1] <- dp
for d in {children of c} do
Dp[0] <- dfs(d, Dp[0])[0]
Dp[1] <- dfs(d, Dp[1])[1]
end for
for i = 0 to X - W[c] do
chmax(Dp[C[c]][i + W[c]], Dp[C[c] ^ 1][i] + B[c])
end for
return Dp
However, this DFS (Depth-First Search) is called for a vertex twice as much time as that for its parent, so the complexity is \(\mathrm{O}^{\ast}(2^N)\).
Let us use the heavy-light decomposition to improve it. Note that we initially assign Dp[0] = Dp[1] = dp
, then call dfs(d, Dp[*])
; thus, we can call the DFS for the first child only once. We can let such a child be the heavy child to obtain the following algorithm:
function dfs(c, dp)
if c is not leaf then
h = (heavy child of c)
Dp <- dfs(h, dp)
for d in {children of c} / {h} do
Dp[0] <- dfs(d, Dp[0])[0]
Dp[1] <- dfs(d, Dp[1])[1]
end for
else
Dp[0] <- dp
Dp[1] <- dp
end if
for i = 0 to X - W[c] do
chmax(Dp[C[c]][i + W[c]], Dp[C[c] ^ 1][i] + B[c])
end for
return Dp
We evaluate the complexity. Suppose that the current subtree has \(n\) vertices, and the subtrees rooted at the children have \(n_1, n_2, \dots\), and \(n_k\) vertices (where \(n_1 \geq n_2 \geq \dots \geq n_k\)). Then, \(f(n)\), defined as the complexity of dfs(c, dp)
, satisfies:
\[f(n) \leq f(n_1) + 2(f(n_2) + \dots + f(n_k)) + \mathrm{O}(X),\]
\[n - 1 = n_1 + n_2 + \dots + n_k.\]
- In the
for d in {children of c} / {h} do
part, one can return values by reference (“move” in C++) to pass a DP table in an \(\mathrm{O}(1)\) time, so the term dependent on \(X\) is \(\mathrm{O}(X)\), not \(\mathrm{O}(kX)\).
Here, since \(f(n)\) is a polynomial of \(n\), so we maximize it by letting the terms for \(n_2, n_3, \dots\) and \(n_k\) be \(n_3=n_4=\dots=n_k=0\). Thus we have
\[f(n) \leq f(n_1) + 2 f(n_2) + \mathrm{O}(X),\]
\[n - 1 = n_1 + n_2, n_1 \geq n_2.\]
Since \(n_1 \geq n_2\), the right hand side is maximized at \(n_1 = n_2 = \dfrac{n-1}{2}\), so
\[f(n) \leq 3 f\left(\frac{n-1}{2}\right) + \mathrm{O}(X);\]
solving this yields
\[f(n) = \mathrm{O}(n^{\log_2{3}} X) = \mathrm{O}(n^{1.59} X).\]
- By the way, regarding the spacial complexity, stepping down to a light child requires reserving a DP table of size \(\mathrm{O}(X)\), so it is \(\mathrm{O}(X \log N)\).
Enumerating \(F(1), F(2), \dots, F(N)\)
Recall that we are originally asked to enumerate the answer for all rooted trees. Calling the DFS function for all vertices costs \(\mathrm{O}(N^{2.59} X)\), which leads to TLE (Time Limit Exceeded), so an optimization is needed.
Observing dfs
function, we realize that dfs(c, dp)
passes the DP table handed from the parent directly to its child. Thus, if we supply the dfs
function with dfs(root, initial dp table)
, function calls with dfs(u, initial dp table)
occur for all vertices \(u\) in the heavy path starting from the root
.
Thus, we can derive an algorithm in which dfs(u, initial dp table)
for all heavy paths, where \(u\) is the vertex on the path nearest to the root.
We consider the complexity. Let \(f(n) = \mathrm{O}(N^{1.59} X)\) be the complexity of the algorithm for a fixed root, and \(g(n)\) be the complexity to enumerate for all roots. If the current subtree has \(n\) vertices and its children’s subtree have \(n_1, n_2, \dots\), and \(n_k\) vertices, then
\[g(n) \leq f(n_1) + g(n_1) + g(n_2) + g(n_3) + \dots + g(n_k).\]
Just as before, it is maximized with \(n_1 = n_2 = (n-1)/2, n_3 = \dots = n_k = 0\), so
\[g(n) \leq 2 g((n-1)/2) + f(n),\]
which yields \(g(n) = \mathrm{O}(N^{1.59} X)\) by solving it.
Therefore, this problem can be solved in a total of \(\mathrm{O}(N^{1.59} X)\) time, which is fast enough.
- Sample code (C++)
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using vl = vector<ll>;
constexpr int nmax = 1010;
#define rep(i, n) for (int i = 0; i < int(n); i++)
#define rep1(i, n) for (int i = 1; i < int(n); i++)
int N, X, W[nmax], C[nmax];
long long B[nmax], ans[nmax];
vector<int> g[nmax];
void chmax(ll& a, ll b) { a = max(a, b); }
array<vl, 2> dfs(int c, const vl& dp, bool memo) {
array<vl, 2> Dp;
if (g[c].empty()) {
Dp[0] = Dp[1] = dp;
} else {
Dp = dfs(g[c][0], dp, memo);
rep1(i, g[c].size()) rep(t, 2) Dp[t] = dfs(g[c][i], Dp[t], false)[t];
}
rep(i, X - W[c] + 1) chmax(Dp[C[c]][i + W[c]], Dp[C[c] ^ 1][i] + B[c]);
if (memo) {
rep(i, X - W[c] + 1) chmax(ans[c], Dp[C[c] ^ 1][i] + B[c]);
rep1(i, g[c].size()) dfs(g[c][i], dp, true);
}
return Dp;
}
int main() {
cin >> N >> X;
vector<int> P(N, -1), sub(N, 1);
for (int c = 1, p; c < N; c++) cin >> p, P[c] = --p, g[p].push_back(c);
for (int i = N - 1; i >= 0; i--) {
if (i) sub[P[i]] += sub[i];
sort(begin(g[i]), end(g[i]), [&](int j, int k) { return sub[j] > sub[k]; });
}
rep(i, N) cin >> B[i] >> W[i] >> C[i];
vl init(X + 1, -2e18);
init[0] = 0;
dfs(0, init, true);
rep(i, N) cout << ans[i] << "\n";
}
posted:
last update: