[CF1388D] Captain Flint and Treasure

Description

有长度为 (n) 的两个数组 (a) 和 (b) 。一开始，答案 (ans) 等于 (0) 。现在我们定义如下操作：选择一个位置 (i) ((1leq i leq n))；让 (ans) 增大 (a_i)；如果 $b_i eq -1 $ 就将 (a_{b_i}) 增大 (a_i)。如果每一个 (i) ((1leq 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] << " ";
    }
}