CF1388D Captain Flint and Treasure

CF1388D Captain Flint and Treasure

大致题意

有两个大小为\(n\)的数组\(a,b\)，进行\(n\)次操作，每次选择一个下标\(i\),使答案\(ans = ans + a[i]\)。若\(b[i] \neq -1\)，则同时使\(a[b[i]] = a[b[i]] + a[i]\)。求一个顺序使得最终答案最大。

\(b[i]\)的指向不会构成环

---

解题思路

前面的操作会影响到后面的结果，所以很容易的想到用类似拓扑排序的思路去做。将\(b[i]\)看做一条从\(i\)连向\(b[i]\)的边。在这个图上进行拓扑排序，不断将当前的\(a[i]\)进行转移。如果当前的\(a[i] > 0\)则使\(a[b[i]] = a[b[i]] + a[i]\)，否则不变。

这样做完之后，我们就得到了答案的最大值，但是这题同时还要求我们输出方案。

考虑到如果对于一个\(a[p] < 0\)，我们不希望它的值传递给它连向的点\(q\),所以我们就需要在选择点\(q\)之后选择点\(p\)。对于需要传递值的\(p\)同理，我们要在选择\(q\)之前选择点\(p\)。这个问题同样可以用拓扑排序解决。我们对两种关系分别建反向边和正向边。在新图上跑一边拓扑排序，就得到了我们需要的顺序方案。

---

代码

#include <bits/stdc++.h>
using namespace std;
using ll       = long long;
const int maxn = 1e6;
int       n;

vector<int> edges[maxn + 10];
int         deg[maxn + 10];
ll          a[maxn + 10];
int         b[maxn + 10];
ll          ans = 0;

vector<int> edges2[maxn + 10];
int         deg2[maxn + 10];

void toposort()
{
    queue<int> q;
    for (int i = 1; i <= n; i++) {
        if (deg[i] == 0)
            q.push(i);
    }
    while (!q.empty()) {
        int p = q.front();
        q.pop();
        ans += a[p];
        for (auto to : edges[p]) {
            deg[to]--;
            if (a[p] > 0) {
                a[to] += a[p];
                edges2[p].emplace_back(to);
                deg2[to]++;
            }
            else {
                edges2[to].emplace_back(p);
                deg2[p]++;
            }
            if (deg[to] == 0)
                q.push(to);
        }
    }
}
vector<int> v;
void        toposort2()
{
    queue<int> q;
    for (int i = 1; i <= n; i++) {
        if (deg2[i] == 0)
            q.push(i);
    }
    while (!q.empty()) {
        int p = q.front();
        q.pop();
        v.emplace_back(p);
        for (auto to : edges2[p]) {
            deg2[to]--;
            if (deg2[to] == 0)
                q.push(to);
        }
    }
}

void solve()
{
    cin >> n;
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
    }
    for (int i = 1; i <= n; i++) {
        cin >> b[i];
        if (b[i] != -1) {
            deg[b[i]]++;
            edges[i].emplace_back(b[i]);
        }
    }
    toposort();
    toposort2();
    cout << ans << "\n";
    for (auto i : v) {
        cout << i << " ";
    }
    cout << "\n";
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    solve();
    return 0;
}