#include "friend.h"
#include <bits/stdc++.h>
using namespace std;
int findSample(int n, int confidence[], int host[], int protocol[]) {
vector<vector<pair<int, int>>> g(n);
for (int i = 1; i < n; ++i) {
g[host[i]].emplace_back(i, protocol[i]);
}
vector dp(n, array<int, 2>{});
//dp[i][0] - don't take neither i nor neighbours of i;
//dp[i][1] - max(dp[i][0], "take some of neighbours, and maybe i")
for (int i = n - 1; i >= 0; --i) {
for (auto [to, type] : g[i]) {
if (type == 0) {
dp[i][0] += dp[to][1];
} else {
dp[i][0] += dp[to][0];
}
}
//take i
dp[i][1] = confidence[i];
for (auto [to, type] : g[i]) {
if (type == 0 || type == 2) {
dp[i][1] += dp[to][0];
} else {
dp[i][1] += dp[to][1];
}
}
//don't take i
vector<int> suf(g[i].size() + 1);
for (int j = int(g[i].size()) - 1; j >= 0; --j) {
suf[j] = suf[j + 1];
auto [to, type] = g[i][j];
if (type == 1 || type == 2) {
suf[j] += dp[to][0];
} else {
suf[j] += dp[to][1];
}
}
int sum0 = 0, sum_good = 0;
for (int j = 0; j < int(g[i].size()); ++j) {
auto [to, type] = g[i][j];
sum0 += dp[to][0];
if (type == 1) {
sum_good += dp[to][1] - dp[to][0];
} else if (type == 2) {
dp[i][1] = max(dp[i][1], suf[j + 1] + sum0 + sum_good + dp[to][1] - dp[to][0]);
}
dp[i][1] = max(dp[i][1], suf[j + 1] + sum0 + sum_good);
}
dp[i][1] = max(dp[i][0], dp[i][1]);
}
return dp[0][1];
}
