Codeforces Round #660 (Div. 2) D. Captain Flint and Treasure

  1 #include<bits/stdc++.h>
  2 using namespace std;
  3 #define ll long long
  4 #define ull unsigned long long
  5 #define pb push_back
  6 const int maxn=2e5+10;
  7 const int mod=1e9+7;
  8 ll a[maxn];
  9 ll b[maxn];
 10 int ru[maxn];
 11 int chu[maxn];
 12 vector<int>E[maxn];
 13 vector<int>G[maxn];
 14 int vis[maxn];
 15 
 16 
 17 
 18 int main()
 19 {
 20     ios::sync_with_stdio(false);
 21     cin.tie(0);
 22     int n;
 23     cin>>n;
 24     ll sum=0;
 25     for(int i=1;i<=n;i++)
 26     {
 27         cin>>a[i];
 28     }
 29     for(int i=1;i<=n;i++)
 30     {
 31         cin>>b[i];
 32     }
 33     vector<int>ans;
 34     for(int i=1;i<=n;i++)
 35     {
 36         if(b[i]==-1)
 37             continue;
 38         ru[b[i]]++;
 39         chu[i]++;
 40         E[i].pb(b[i]);
 41         G[b[i]].pb(i);
 42     }
 43     queue<int>q;
 44     for(int i=1;i<=n;i++)
 45     {
 46         if(ru[i]==0)
 47         {
 48             q.push(i);
 49         }
 50     }
 51     //cout<<sum<<‘\n‘;
 52     while(!q.empty())
 53     {
 54         int u=q.front();
 55         q.pop();
 56         for(auto &v:E[u])
 57         {
 58             ru[v]--;
 59             if(a[u]>=0)
 60             {
 61                 a[v]+=a[u];
 62                 sum+=a[u];
 63                 ans.pb(u);
 64                 vis[u]=1;
 65             }
 66             if(ru[v]==0)
 67             {
 68                 q.push(v);
 69             }
 70         }
 71     }
 72     //cout<<sum<<‘\n‘;
 73     queue<int>qh;
 74     for(int i=1;i<=n;i++)
 75     {
 76         if(chu[i]==0&&vis[i]==0)
 77         {
 78             qh.push(i);
 79         }
 80     }
 81     //cout<<sum<<‘\n‘;
 82     while(!qh.empty())
 83     {
 84         int u=qh.front();
 85         qh.pop();
 86         if(vis[u]==0)
 87         {
 88             sum+=a[u];
 89             ans.pb(u);
 90         }
 91         for(auto &v:G[u])
 92         {
 93             chu[v]--;
 94             if(chu[v]==0)
 95                 qh.push(v);
 96         }
 97     }
 98     cout<<sum<<‘\n‘;
 99     for(auto &v:ans)
100     {
101         cout<<v<<" ";
102     }
103 
104 
105 
106 
107 }