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] << " ";
    }
}