#include "friend.h"
#include <bits/stdc++.h>
using namespace std;
vector<vector<int>> adj;
int findSample(int n, int confidence[], int host[], int protocol[]){
int mip = 4, map = 0;
for(int i = 1; i < n; i++) {
mip = min(mip, protocol[i]);
map = max(map, protocol[i]);
}
if(n <= 10) { // Subtask 1
adj.resize(n);
for(int i = 1; i < n; i++) {
int x, y;
if(protocol[i] == 0) {
tie(x, y) = {i, host[i]};
adj[x].push_back(y);
adj[y].push_back(x);
}
if(protocol[i] == 1) {
for(int g : adj[host[i]]) {
adj[i].push_back(g);
adj[g].push_back(i);
}
}
if(protocol[i] == 2) {
for(int g : adj[host[i]]) {
adj[i].push_back(g);
adj[g].push_back(i);
}
tie(x, y) = {i, host[i]};
adj[x].push_back(y);
adj[y].push_back(x);
}
}
int best = 0;
for(int i = 0; i < (1 << n); i++) {
vector<int> goods;
for(int j = 0; j < n; j++) {
if(i & (1 << j)) goods.push_back(j);
}
bool ok = true;
for(int g : goods) {
for(int h : goods) {
if(g == h) continue;
for(int k : adj[g]) {
if(k == h) ok = false;
}
}
}
if(!ok) continue;
int ans = 0;
for(int g : goods) {
ans += confidence[g];
}
best = max(best, ans);
vector<int>().swap(goods);
}
return best;
}
else if(mip == map && mip == 1) { // Subtask 2
int ans = 0;
for(int i = 0; i < n; i++) ans += confidence[i];
return ans;
}
else if(mip == map && mip == 2) { // Subtask 3
sort(confidence, confidence+n);
return confidence[n-1];
}
else if(mip == map) { // Subtask 4
// It is a tree
adj.resize(n);
int ans[1001][2];
ans[0][0] = ans[0][1] = 0;
for(int i = 1; i < n; i++) {
ans[i][0] = ans[i][1] = 0;
adj[i].push_back(host[i]);
adj[host[i]].push_back(i);
}
function<int(int, int, int)> dfs = [&](int n, int x, int par) {
if(ans[n][x]) return ans[n][x];
if(x == 0) {
int sum = 0;
for(int ch : adj[n]) {
if(ch == par) continue;
sum += dfs(ch, 1, n);
}
return ans[n][x] = sum;
}
int sum = confidence[n];
for(int ch : adj[n]) {
if(ch == par) continue;
sum += dfs(ch, 0, n);
}
return ans[n][x] = max(sum, dfs(n, 0, par));
};
return dfs(0, 1, -1);
}
else if(map - mip == 1) { // Subtask 5
int col[1001];
col[0] = 0;
for(int i = 1; i < n; i++) {
col[i] = (col[host[i]]) ^ (1 ^ (protocol[i]));
}
int x = 0;
for(int i = 0; i < n; i++) {
if(col[i]) x++;
}
return max(x, n-x);
}
// Subtask 6
return -1;
}
