#include "friend.h"
#include <bits/stdc++.h>
using namespace std;
pair<int, int> dfs(int v, vector<int> adj[], int val[])
{
int pick = val[v], no_pick = 0;
for(int u : adj[v])
{
auto [x1, x2] = dfs(u, adj, val);
pick += x2;
no_pick += x1;
}
return make_pair(max(pick, no_pick), no_pick);
}
int findSample(int n,int confidence[],int host[],int protocol[]){
vector<int> g_main[n], g_group[n];
int head[n];
int tot[n];
vector<int> con = {0};
for(int i = 0; i < n; i++)
head[i] = i, tot[i] = 0;
for(int i = 1; i < n; i++)
{
int par = host[i];
int par_h = head[par];
if(protocol[i] == 0)
{
g_main[par_h].push_back(i);
con.push_back(i);
}
else
{
head[i] = par_h;
if(protocol[i] == 2)
{
g_group[par].push_back(i);
}
else
{
con.push_back(i);
}
}
}
for(int v : con)
{
auto [ans, _] = dfs(v, g_group, confidence);
//cout << " head " << v << " ans : " << ans << endl;
tot[head[v]] += ans;
}
auto [ans, _] = dfs(0, g_main, tot);
return ans;
}
