Official

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] and Dp[1] are used as temporary variables that mean:
    • c: currently inspected vertex
    • Dp[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 modifying dp, it is sufficient to pass a const 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 at c.”
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: