[CF1388D] Captain Flint and Treasure - 树形dp

Description

有长度为 \(n\) 的两个数组 \(a\) 和 \(b\) 。一开始，答案 \(ans\) 等于 \(0\) 。现在我们定义如下操作：选择一个位置 \(i\) \((1\leq i \leq n)\)；让 \(ans\) 增大 \(a_i\)；如果 $b_i \neq -1 $ 就将 \(a_{b_i}\) 增大 \(a_i\)。如果每一个 \(i\) \((1\leq i \leq n)\) 只能选一次，请问 \(ans\) 最大是多少？ 并给出 \(ans\) 最大时选择 \(i\) 的顺序。

Solution

对每个顶点算一个累积值 \(f[i]\)，如果 \(f[i]>0\) 就向上更新。

最后扫一遍，对于一个点的所有孩子，先走那些 \(f[i]>0\) 的，再走根，最后走那些 \(f[i]<0\) 的。

#include <bits/stdc++.h>
using namespace std;

#define int long long
const int N = 1000005;

int n, a[N], b[N], f[N];
vector<int> g[N];

signed main()
{
    ios::sync_with_stdio(false);

    cin >> n;
    for (int i = 1; i <= n; i++)
        cin >> a[i];
    for (int i = 1; i <= n; i++)
        cin >> b[i];

    for (int i = 1; i <= n; i++)
    {
        if (b[i] != -1)
        {
            g[b[i]].push_back(i);
        }
    }

    int ans = 0;

    function<void(int)> dfs = [&](int p) -> void {
        f[p] = a[p];
        for (int q : g[p])
        {
            dfs(q);
            if (f[q] > 0)
            {
                f[p] += f[q];
            }
        }
        ans += f[p];
    };

    for(int i=1;i<=n;i++) if(b[i]==-1) dfs(i);

    cout << ans << endl;

    vector<int> output;

    function<void(int)> gen = [&](int p) -> void {
        for (int q : g[p])
        {
            if (f[q] > 0)
            {
                gen(q);
            }
        }

        output.push_back(p);

        for (int q : g[p])
        {
            if (f[q] <= 0)
            {
                gen(q);
            }
        }
    };
    for(int i=1;i<=n;i++) if(b[i]==-1) gen(i);


    for (int i = 0; i < output.size(); i++)
    {
        cout << output[i] << " ";
    }
}