#include "friend.h"
#include<bits/stdc++.h>
using namespace std;
struct UnionFind{
vector<int> par;
UnionFind(int n){
par.resize(n);
iota(par.begin(),par.end(),0);
}
int find(int i){
if(par[i] == i) return i;
return par[i] = find(par[i]);
}
void unite(int i, int j){
i = find(i);
j = find(j);
par[j] = i;
}
};
vector<int> arr,dpT, dpN,add;
vector<vector<int>> adj;
void DFS(int v){
dpT[v] = arr[v]+add[v];
int oth1 = 0;
for(auto u : adj[v]){
DFS(u);
}
for(auto u : adj[v]){
dpT[v] += dpN[u];
dpN[v] += max(dpT[u]-add[u],dpN[u])+add[u];
}
}
// Find out best sample
int findSample(int n,int confidence[],int host[],int protocol[]){
int ans=0;
adj.resize(n);
arr.resize(n);
add.resize(n);
vector<vector<int>> additions(n);
arr[0] = confidence[0];
UnionFind uf(n);
stack<int> roots;
roots.push(0);
for(int i = 1; i < n; i++){
arr[i] = confidence[i];
if(protocol[i] == 0){
adj[host[i]].push_back(i);
}else if(protocol[i] == 1){
host[i] = uf.find(host[i]);
additions[host[i]].push_back(i);
roots.push(i);
}else{
uf.unite(host[i],i);
int real =uf.find(i);
arr[real] = max(arr[real],confidence[i]);
}
}
dpN = dpT = vector<int>(n);
while(roots.size()){
int r =roots.top(); roots.pop();
DFS(r);
if(r == 0) break;
add[host[r]] += max(dpT[r], dpN[r]);
}
for(int i = 0; i < n; i++){
ans = max({ans, dpT[i], dpN[i]});
}
return ans;
}
