#include <bits/stdc++.h>
#include "tree.h"
using namespace std;
typedef long long ll;
int n;
vector<int> p, w;
vector<vector<ll>> adj;
map<ll, ll> leaves; // leaf count, tree count
ll zeroCnt = 0;
ll dfs(ll cur) {
if (w[cur] == 0 && p[cur] >= 0 && w[p[cur]] == 1) zeroCnt++;
ll leafCnt = 0;
for (auto &e : adj[cur]) {
leafCnt += dfs(e);
}
if (w[cur] == 1 && (p[cur] == -1 || w[p[cur]] == 0)) {
leaves[max(1ll, w[cur] ? leafCnt : 1)]++;
}
return max(1ll, w[cur] ? leafCnt : 1);
}
void init(vector<int> P, vector<int> W) {
p = P;
w = W;
n = (int)p.size();
adj = vector<vector<ll>>(n);
zeroCnt = 0;
for (int i = 1; i < n; i++) {
adj[p[i]].push_back(i);
}
dfs(0);
}
ll query(int L, int R) {
ll result = -zeroCnt * L;
for (auto &[leafCnt, c] : leaves) {
ll leafStuff = leafCnt * L;
if (leafStuff > R) {
result += c * (2 * leafStuff - R);
}
else {
result += c * leafStuff;
}
}
return result;
}
#ifdef TEST
#include "grader.cpp"
#endif
