公式

Ex - Many Illumination Plans 解説 by Nyaan


以下ではまず \(v=1\) の場合の答え \(F(1)\) を計算する問題を考えます。

自然な DP, \(\mathrm{O}(N X^2)\)

まず、自然な木 DP を実装してみましょう。すると次のようになり、\(\mathrm{O}(N X^2)\) で計算できるアルゴリズムを構成できます。

  • 以降の疑似コードでは c, Dp[0], Dp[1] は次の意味を持つ一時変数として利用されています。
    • c : 現在見ている頂点
    • Dp[0], Dp[1] : 長さ \(X+1\) の DP テーブル
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

このアルゴリズムでは max-plus convolution の部分が \(\mathrm{O}(X^2)\) かかってしまい全体で \(\mathrm{O}(NX^2)\) の計算量となりますが、\(X\) が非常に大きい制約下で AC するのは困難でしょう。

重軽再帰 DP (HLRecDP), \(\mathrm{O}(N^{1.59} X)\)

以降の解説は tmaehara 氏による記事, 論文 を参考にしました。

この問題は 重軽再帰 DP (HLRecDP) と呼ばれるテクニックを利用することで \(\mathrm{O}(\mathrm{poly}(N) X)\) で解くことができます。

\(\mathrm{O}(NX^2)\) の DP は部分木ごとに DP テーブルをボトムアップにマージしていくアルゴリズムでしたが、マージに \(\mathrm{O}(X^2)\) かかるという問題点がありました。
では、逆にマージを発生させないようにトップダウンに DP を計算することを考えてみましょう。すると、次の再帰関数で dfs(root, initial dp table) を呼び出すことで根が root であるときの DP テーブルの値を計算できるとわかります。(これは慣れない DP なので少し頭を使います。人によっては DSU on Tree アルゴリズムの気持ちで考えると理解できるかもしれません)

  • 注 1 : dfs(c, dp)dp参照渡し です。(ただし dp 自体を更新することはないので C++ では const reference で渡せばよい)
  • 注 2 : ボトムアップ型の DP と違って、途中の dfs(c, dp) の返り値は 「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

ただしこの DFS は1 つ深い頂点に進むごとに呼び出しが 2 倍になっているので計算量は \(\mathrm{O}^{\ast}(2^N)\) になっています。

重軽分解を利用して工夫してみましょう。はじめ Dp[0] = Dp[1] = dp を代入して、その後 dfs(d, Dp[*]) を呼んでいるのに注目すると、DFS で最初に訪問する子は DFS 呼び出しを 1 回に済ますことができます。このような子に heavy child を割り当てると次のアルゴリズムになります。

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

計算量を考えます。今見ている部分木の頂点数が \(n\) で、子の部分木の頂点数が \(n_1, n_2, \dots, n_k\) であるとします。 (ただし \(n_1 \geq n_2 \geq \dots \geq n_k\)) このとき、dfs(c, dp) の計算量を \(f(n)\) とすると

\[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\]

という式が成り立ちます。

  • for d in {children of c} / {h} do の部分では返り値を参照渡し (C++ でのムーブ) すれば \(\mathrm{O}(1)\) で DP テーブルを渡せるので、 \(X\) に依存する部分は \(\mathrm{O}(kX)\) ではなく \(\mathrm{O}(X)\) になります。

ここで、\(f(n)\)\(n\) に関する多項式なので \(n_2, n_3, \dots, n_k\) の部分については \(n_3=n_4=\dots=n_k=0\) にするのが最大です。よって

\[f(n) \leq f(n_1) + 2 f(n_2) + \mathrm{O}(X)\]

\[n - 1 = n_1 + n_2, n_1 \geq n_2\]

という式が成り立ちます。右辺は \(n_1 \geq n_2\) を踏まえると \(n_1 = n_2 = \dfrac{n-1}{2}\) のときが最大なので

\[f(n) \leq 3 f\left(\frac{n-1}{2}\right) + \mathrm{O}(X)\]

となり、これを解いて

\[f(n) = \mathrm{O}(n^{\log_2{3}} X) = \mathrm{O}(n^{1.59} X)\]

を得られます。

  • 余談ですが、空間計算量については、 light child に潜るごとに新たに \(\mathrm{O}(X)\) の DP テーブルが確保されているので \(\mathrm{O}(X \log N)\) であるとわかります。

\(F(1), F(2), \dots, F(N)\) の列挙

さて、この問題は全ての根付き木に対して答えを列挙する問題でした。愚直に全ての頂点に対して dfs 関数を呼び出すと \(\mathrm{O}(N^{2.59} X)\) かかり TLE するので計算量改善が必要です。
dfs 関数を観察すると、dfs(c, dp) は heavy child に親から来る DP テーブルをそのまま子に渡す挙動をしているのがわかります。よって、dfs 関数に dfs(root, initial dp table) を渡すと、root から heavy path で結ばれた頂点列に含まれる頂点 \(u\) に対して dfs(u, initial dp table) が呼び出されることがわかります。

ここから、全ての heavy path に対して、heavy path の根側の頂点を \(u\) として、 dfs(u, initial dp table) を呼び出すアルゴリズムが導出できます。

計算量を考えます。根を固定したときのアルゴリズムの計算量を \(f(n) = \mathrm{O}(N^{1.59} X)\), 全ての根に対して列挙するときの計算量を \(g(n)\) とします。今見ている部分木の頂点数が \(n\) で、子の部分木の頂点数が \(n_1, n_2, \dots, n_k\) ですると

\[g(n) \leq f(n_1) + g(n_1) + g(n_2) + g(n_3) + \dots + g(n_k)\]

になり、前と同様に \(n_1 = n_2 = (n-1)/2, n_3 = \dots = n_k = 0\) とするのが最大なので

\[g(n) \leq 2 g((n-1)/2) + f(n)\]

となり、この式を解いて \(g(n) = \mathrm{O}(N^{1.59} X)\) を得られます。
以上よりこの問題を \(\mathrm{O}(N^{1.59} X)\) で解くことができて、これは十分高速です。

  • 実装例(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";
}

投稿日時:
最終更新: