#include "friend.h"
#include <algorithm>
#include <vector>
struct Edge
{
Edge(int e, bool d, bool b) : end(e), isDirect(d), isBlocking(b) {}
int end;
bool isDirect, isBlocking;
};
const int MAXN = 1e5;
std::vector<Edge> network[MAXN];
int value[MAXN];
int bestSample[MAXN][2][2];
void sample(int node)
{
int sumBlocking = 0;
for (Edge &edge : network[node])
{
sample(edge.end);
int weight = std::max(std::max(bestSample[edge.end][0][0], bestSample[edge.end][0][1]), std::max(bestSample[edge.end][1][0], bestSample[edge.end][1][1]));
if (edge.isBlocking)
{
bestSample[node][0][0] += bestSample[edge.end][0][0];
bestSample[node][1][0] += bestSample[edge.end][0][0];
if (edge.isDirect)
{
bestSample[node][1][1] += bestSample[edge.end][0][0];
bestSample[node][0][1] = std::max(bestSample[node][0][1] + bestSample[edge.end][0][0], sumBlocking + weight);
sumBlocking += bestSample[edge.end][0][0];
}
else
{
bestSample[node][1][1] += weight;
sumBlocking += weight;
bestSample[node][0][1] = std::max(bestSample[node][0][1] + bestSample[edge.end][0][0], sumBlocking);
}
}
else
{
bestSample[node][0][0] += weight;
bestSample[node][1][0] += bestSample[edge.end][0][0];
bestSample[node][1][1] += bestSample[edge.end][0][0];
sumBlocking += bestSample[edge.end][0][0];
bestSample[node][0][1] += weight;
}
}
}
int findSample(int n, int confidence[], int host[], int protocol[])
{
for (int iNode = 0; iNode < n; iNode++)
bestSample[iNode][1][0] = bestSample[iNode][1][1] = confidence[iNode];
for (int iEdge = 1; iEdge < n; iEdge++)
network[host[iEdge]].emplace_back(iEdge, (protocol[iEdge] % 2 == 0), (protocol[iEdge] > 0));
sample(0);
return std::max(std::max(bestSample[0][0][0], bestSample[0][0][1]), std::max(bestSample[0][1][0], bestSample[0][1][1]));
}
