優秀なプログラマになるには (AOJ1631)

この記事は帰ってきた AOJ-ICPC Advent Calendar 2022 10 日目の記事です。

少し珍しいタイプの木DPの問題です。

問題概要

 N 頂点の木がある。頂点  i の親は  P_i である。各頂点  i には価値  S_i のコインが  H_i 枚ある。 あなたは木から価値の合計が最大になるよう最大  K 枚のコインを回収したい。ただし、頂点  i のコインを 1 枚以上回収するためには、 P_i のコインを  L_i 枚以上回収していなければならない。

条件を満たすようにコインを回収するとき、価値の合計の最大値はいくつか。

制約

  •  N \le 100
  •  K, H_i \le 10 ^ 5
  •  S_i \le 10 ^ 9

解法

問題設定と制約から、まずは単純な木 DP(dp[i][j] := (i 以下の部分木で j 枚回収するときの最大値)) を考えてしまった。しかし、これはなかなか筋が悪いことがわかる。 というのも、この問題の肝の設定である。( P_i のコインを  L_i 枚以上回収していなければならない。) という条件に相性が悪い。この条件を上手く省けるような DP を設計したい。

条件を活かすためには、下からではなく上から dfs 順に見るいわゆるトップダウンの木 DP を考えると、今訪問している頂点までの条件を満たしていることが確定しているので、一見相性がよさそうに見える。 しかし、このパターンでは先祖でも子孫でもない頂点の情報が完全に抜け落ちてしまっているため、DP するのは難しそうに見える。

この情報は、実は DP する順番を上手く定めることで保持することが出来る。 トップダウンの木 DP では、DP テーブルを参照として持たせながら dfs をするパターンと、更新を子以外も含めることで上手く行う二つのパターンが存在する。 今回は後者を用いて実装した。(前者の方がわかりやすいが、この問題は定数倍がきつかったのでやむなく)

子以外を含めるというのは、例えばある葉に到達したときに、その葉をもとに更新する dp テーブルを dfs 順で次のノードにするということである。 こうすることで、全ての頂点を訪問した情報を得ることが出来る。

さて、情報が抜け落ちてしまうという問題に戻ろう。 直接の子供以外に遷移するとき、それは先祖の兄弟の関係にあるノードである。もしこのノードの  L_i の値が既に訪問している兄弟ノードの  L_i よりも小さければ、条件は無視できる。

この考察がこの問題の本質である。

よって、条件は無視できるので単に更新することが出来る。

更新は以下のようにして行われる。

頂点  x を見ているとき

子供  v に対して、dp[v] を dp[x] から  L_v 枚以上  H_x 枚以下を回収した上で更新。 先祖の兄弟  v に対して、dp[v] を dp[x] から  0 枚以上  H_x 枚以下を回収した上で更新。

どの子供にも遷移しない場合があるので、葉以外の頂点からも dfs 順で最も近い先祖の弟に遷移させなければならないことに注意する(ぼくはこれに気づかなくて死んでました)

この更新、パッと見では  \mathrm{O}(K ^ 2) に見えるが、実はこれは簡単に  \mathrm{O}(K) に落とすことが出来る。

dp[v][i] =  \underset{i - H_x \le j \le i - L_v}{\min} \left( dp_{x, j} + (i - j) \times S_x \right)

 = \underset{i - H_x \le j \le i - L_v}{\min} \left( dp_{x, j} - j \times S_x \right)+ i \times S_x

これはスライド最小値で計算できる。

実装

#ifdef noimi
#include <stdc++.h>
#else
#include <bits/stdc++.h>
#endif

#define ll long long

#define REP0(n) for(ll iiii = 0; iiii < (n); i++)
#define REP1(i, n) for(ll i = 0; i < (n); i++)
#define REP2(i, a, b) for(ll i = (a); i < (b); i++)
#define overload3(a, b, c, name, ...) name
#define rep(...) overload3(__VA_ARGS__, REP2, REP1, REP0)(__VA_ARGS__)

#define vi vector<int>
using namespace std;

template <typename T, typename S> bool chmax(T &x, const S &y) { return (x < y ? x = y, true : false); }

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
    os << "{";
    rep(i, v.size()) {
        if(i) os << ", ";
        os << v[i];
    }
    return os << "}";
}

// deque
struct DEQUE {
    int l, r;
    int d[100100];
    DEQUE() : l(0), r(0) {}
    bool empty() { return l == r; }
    void push_back(int x) { d[r++] = x; }
    void pop_back() { r--; }
    void pop_front() { l++; }
    void clear() { l = r = 0; }
    int back() { return d[r - 1]; }
    int front() { return d[l]; }
};

ll dp[101][100010];

int main() {
    DEQUE d;
    while(true) {
        int n, k;
        cin >> n >> k;
        if(!n) exit(0);

        vi h(n);
        vector<ll> s(n);
        vector<vi> g(n);
        rep(i, n) cin >> h[i], h[i] = min(h[i], k);
        rep(i, n) cin >> s[i];

        rep(i, 1, n) {
            int x;
            cin >> x;
            x--;
            g[x].emplace_back(i);
        }
        vi l(n);
        rep(i, 1, n) cin >> l[i];

        constexpr ll inf = 1e18;
        rep(i, n + 1) rep(j, k + 1) dp[i][j] = -inf;

        vector<int> idx, nxt(n);
        {
            auto dfs = [&](auto &&f, int x, int nv) -> void {
                idx.emplace_back(x);
                sort(begin(g[x]), end(g[x]), [&](int i, int j) { return l[i] > l[j]; });
                rep(i, g[x].size()) { f(f, g[x][i], (i == g[x].size() - 1 ? nv : g[x][i + 1])); }
                nxt[x] = nv;
            };
            dfs(dfs, 0, n);
        }

        dp[0][0] = 0;

        auto f = [&](ll *tmp, ll *nxt, ll s, ll L, ll R) {
            // rep(i, k + 1) tmp[i] -= i * s;
            d.clear();
            rep(i, L, k + 1) {
                while(!d.empty() and tmp[d.back()] - d.back() * s < tmp[i - L] - (i - L) * s) d.pop_back();
                d.push_back(i - L);
                chmax(nxt[i], tmp[d.front()] - d.front() * s + i * s);
                if(d.front() == i - R) d.pop_front();
            }
        };

        rep(i, n) {
            int x = idx[i];
            for(auto e : g[x]) {
                int L = l[e], R = h[x];
                f(dp[x], dp[e], s[x], L, R);
            }
            int e = idx[i + 1];
            f(dp[x], dp[nxt[x]], s[x], 0, h[x]);
        }

        ll ma = 0;
        rep(j, k + 1) chmax(ma, dp[n][j]);
        cout << ma << endl;
    }
}