#include "tree.h"
#include <stdio.h>
#include <map>
#include <set>
#include <queue>
#define lol long long
using namespace std;
int n;
vector<int> p, w;
vector<vector<lol> > children;
map<lol, lol> counts;
lol zeroes;
// mode == 0: W <= 1
// mode == 1: W[P[i]] <= W[i] for each 1 <= i <= N
lol mode;
vector<int> initialGrades;
// return countLeaves
lol dfs(lol node)
{
lol leaves = 0;
if (w[node] == 0) {
for (auto j: children[node]) {
lol result = dfs(j);
counts[result]++;
}
return 1;
}
for (auto j: children[node])
leaves += dfs(j);
return (children[node].size() == 0) ? 1 : leaves;
}
void init(vector<int> P, vector<int> W) {
p = P;
w = W;
n = (int)p.size();
counts.clear();
children.clear();
mode = 0;
for (auto j: W) {
if (j != 0 && j != 1) {
mode = 1;
break;
}
}
children = vector<vector<lol> >(p.size(), vector<lol>());
for (int i = 0; i < n; i++) {
if (p[i] != -1)
children[p[i]].push_back(i);
zeroes += w[i] == 0;
}
lol leavesOfRoot = dfs(0);
counts[leavesOfRoot]++;
initialGrades = vector<int>(p.size(), 0);
for (int i = 0; i < n; i++) {
if (p[i] == -1)
continue;
initialGrades[p[i]]++;
}
}
lol subtree(lol L, lol R, lol leavesCount) {
lol sumLeaves = leavesCount * ((lol) L);
if (sumLeaves <= R) {
return sumLeaves;
} else {
return sumLeaves + (sumLeaves - R);
}
}
long long query(int L, int R) {
lol result = 0;
if (mode == 0) {
for (auto j: counts)
result += subtree(L, R, j.first) * j.second;
return result - zeroes * L;
} else if (mode == 1) {
vector<int> currentGrades = vector<int>(n, 0);
vector<lol> currentValues = vector<lol>(n, 0);
queue<int> q;
for (int i = 0; i < n; i++) {
currentGrades[i] = initialGrades[i];
if (currentGrades[i] == 0)
q.push(i);
}
while (!q.empty()) {
int i = q.front();
q.pop();
if (currentValues[i] < L) {
result += w[i] * (L - currentValues[i]);
currentValues[i] = L;
} else if (currentValues[i] > R) {
result += w[i] * (currentValues[i] - R);
currentValues[i] = R;
}
if (i == 0)
continue;
currentValues[p[i]] += currentValues[i];
currentGrades[p[i]]--;
if (currentGrades[p[i]] == 0)
q.push(p[i]);
}
return result;
}
return -1;
}
