#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
int cs = children[i].size();
for (int c : children[i]) {
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)
auto contribution_nn = [&] (int j, int midpoint) {
if (midpoint >= cs) return 0;
if ((j < midpoint && connects_self(children[i][j])) ||
(j > midpoint && connects_neighbors(children[i][j]))) {
return best_yn[children[i][j]];
} else {
return max(best_nn[children[i][j]], best_ny[children[i][j]]);
}
};
if (cs > 0) {
int current_nn = 0;
for (int j = 0; j < cs; ++j) current_nn += contribution_nn(j, 0);
best_nn[i] = current_nn;
for (int midpoint = 1; midpoint <= cs; ++midpoint) {
current_nn -= contribution_nn(midpoint-1, midpoint-1) + contribution_nn(midpoint, midpoint-1);
current_nn += contribution_nn(midpoint-1, midpoint) + contribution_nn(midpoint, midpoint);
best_nn[i] = max(best_nn[i], current_nn);
}
}
auto contribution_ny = [&] (int j, int midpoint) {
if (midpoint >= cs) return 0;
if ((j < midpoint && connects_self(children[i][j])) ||
(j > midpoint && connects_neighbors(children[i][j]))) {
return best_yn[children[i][j]];
} else if (connects_self(children[i][j])) {
return best_yn[children[i][j]];
} else {
return max(best_nn[children[i][j]], best_ny[children[i][j]]);
}
};
// similarly, compute best_ny
if (cs > 0) {
int current_ny = 0;
for (int j = 0; j < cs; ++j) current_ny += contribution_ny(j, 0);
best_ny[i] = current_ny;
for (int midpoint = 1; midpoint <= cs; ++midpoint) {
current_ny -= contribution_ny(midpoint-1, midpoint-1) + contribution_ny(midpoint, midpoint-1);
current_ny += contribution_ny(midpoint-1, midpoint) + contribution_ny(midpoint, midpoint);
best_ny[i] = max(best_ny[i], current_ny);
}
}
best_ny[i] += confidence[i];
}
return max(best_nn[0], best_ny[0]);
}
