#include <bits/stdc++.h>
#include "friend.h"
using namespace std;
// Find out best sample
int findSample(int n,int confidence[],int host[],int protocol[]){
auto connects_self = [&] (int v) {
return protocol[v] == 0 || protocol[v] == 2;
};
auto connects_neighbors = [&] (int v) {
return protocol[v] == 1 || protocol[v] == 2;
};
vector<vector<int>> children(n, vector<int>());
for (int i = 1; i < n; ++i) children[host[i]].push_back(i);
// first y/n: does there exist an indirect neighbor we picked
// (path of neighbor-connects that goes up, then a direct connect,
// then the path goes back down through neighbor-connects)
// second y/n: did we pick this node?
// note that yy cannot happen; since the indirect neighbor and this node
// are adjacent
vector<int> best_nn(n), best_yn(n), best_ny(n);
for (int i = n - 1; i >= 0; --i) {
// best_yn: NOT ALLOWED to take best_n* on any node that connects_neighbors.
// on remaining nodes, we ARE allowed to take best_n*
//
// best_ny: NOT ALLOWED to take best_ny on any node that connects_self.
// remaining nodes should follow restrictions below
//
// best_nn: no restrictions otherwise
// restriction on nodes that take n*:
// - from left to right, must be (n only) ... (ns)? ... (s only)
// all other nodes must be marked yn
for (int c : children[i]) {
int cs = children[i].size();
if (connects_neighbors(c)) {
best_yn[i] += best_yn[c];
} else {
best_yn[i] += max(best_ny[c], best_nn[c]);
}
// midpoint is the latest possible n OR earliest possible s (or both)
for (int midpoint = 0; midpoint <= cs; ++midpoint) {
int current_value = 0;
for (int j = 0; j < cs; ++j) {
if ((j < midpoint && connects_self(children[i][j])) ||
(j > midpoint && connects_neighbors(children[i][j]))) {
current_value += best_yn[children[i][j]];
} else {
current_value += max(best_nn[children[i][j]], best_ny[children[i][j]]);
}
}
best_nn[i] = max(best_nn[i], current_value);
}
// similarly, compute best_ny
for (int midpoint = 0; midpoint <= cs; ++midpoint) {
int current_value = 0;
for (int j = 0; j < cs; ++j) {
if ((j < midpoint && connects_self(children[i][j])) ||
(j > midpoint && connects_neighbors(children[i][j]))) {
current_value += best_yn[children[i][j]];
} else if (connects_self(children[i][j])) {
current_value += best_yn[children[i][j]];
} else {
current_value += max(best_nn[children[i][j]], best_ny[children[i][j]]);
}
}
best_ny[i] = max(best_ny[i], current_value);
}
}
best_ny[i] += confidence[i];
}
return max(best_nn[0], best_ny[0]);
}
