#include "tree.h"
#include <cassert>
#include <set>
#include <iostream>
#define int long long
using i32 = int32_t;
using namespace std;
struct MyVec {
vector<int> vals, suff_sum; int N;
MyVec(const vector<int> &_v) : vals(_v), N(_v.size()) {
sort(vals.begin(), vals.end());
suff_sum.resize(vals.size() + 1);
for (int i = N-1; i >= 0; --i) {
suff_sum[i] = vals[i] + suff_sum[i+1];
}
}
/// returns {sum, count} of values >= target
pair<int, int> sum_geq(int target) {
int i = lower_bound(vals.begin(), vals.end(), target) - vals.begin();
return {suff_sum[i], N - i};
}
};
int N;
vector<int> parent, weight;
vector<vector<int>> children;
namespace St45 {
bool active = true; // set to false in init
int sum_weight_real_leaves = 0;
MyVec tool({});
void prepare() {
vector<bool> seen(N);
vector<int> sub_leaves;
for (int subroot = 0; subroot < N; ++subroot) {
if (children[subroot].empty())
sum_weight_real_leaves += weight[subroot];
if (weight[subroot] > 0 && !seen[subroot]) {
sub_leaves.push_back(0);
vector<int> stk { subroot };
while (!stk.empty()) {
int cur = stk.back();
stk.pop_back();
seen[cur] = true;
if (weight[cur] == 0 || children[cur].empty()) {
++sub_leaves.back();
} else {
for (int child : children[cur]) {
stk.push_back(child);
}
}
}
}
}
tool = MyVec(sub_leaves);
}
int solve(int L, int R) {
// int ans = L*real_leaves;
// for (int x : tool.vals)
// ans += max(L*x - R, 0ll);
// return ans;
auto [sum_above, cnt_above] = tool.sum_geq((R+L-1)/L); // ?
return L*sum_weight_real_leaves + L*sum_above - R*cnt_above;
}
}
void init(vector<i32> _P, vector<i32> _W) {
parent = vector<int>(_P.begin(), _P.end());
weight = vector<int>(_W.begin(), _W.end());
N = (int)parent.size();
children.resize(N);
for (int node = 1; node < N; ++node) {
children[parent[node]].push_back(node);
}
for (int w : weight) {
if (w > 1) St45::active = false;
}
if (St45::active) St45::prepare();
}
struct Helper {
int cost, stock;
auto operator<=>(const Helper&) const = default;
};
class MyPq {
friend MyPq* merge(MyPq*, MyPq*);
public:
void clear() {
helpers.clear();
sum_stock = 0;
}
void push(int cost, int stock) {
if (stock == 0) return;
helpers.insert({cost, stock});
sum_stock += stock;
}
void trim(int lim_sum) {
if (sum_stock > lim_sum) {
consume(sum_stock - lim_sum, true);
assert(sum_stock == lim_sum);
}
}
int cheap_help(int target) {
return consume(target, false);
}
private:
int consume(int target, bool worst) {
int cur_cost = 0;
while (target > 0) {
assert(!helpers.empty());
auto it = worst ? prev(helpers.end()) : helpers.begin();
Helper nxt = *it;
helpers.erase(it);
int taking = min(target, nxt.stock);
target -= taking;
sum_stock -= taking;
nxt.stock -= taking;
cur_cost += taking * nxt.cost;
if (nxt.stock)
helpers.insert(nxt);
}
return cur_cost;
}
multiset<Helper> helpers;
int sum_stock = 0;
};
MyPq* merge(MyPq* small, MyPq* large) {
if (small->helpers.size() > large->helpers.size())
swap(small, large);
for (const Helper &h : small->helpers)
large->push(h.cost, h.stock);
small->clear();
return large;
}
int solve(int L, int R) {
vector<MyPq> block_allocation(N);
vector<MyPq*> dp(N);
/// not always up to date!
vector<int> tmp_sum_subtree(N);
for (int i = 0; i < N; ++i) {
dp[i] = &block_allocation[i];
}
int answer = 0;
for (int node = N-1; node >= 0; --node) {
int &cur_sum = tmp_sum_subtree[node];
if (children[node].empty()) {
answer += weight[node] * L;
cur_sum = L;
continue;
}
MyPq* &cur = dp[node];
for (int child : children[node]) {
cur = merge(cur, dp[child]);
dp[child] = nullptr;
cur_sum += tmp_sum_subtree[child];
}
cur->push(weight[node], cur_sum - L);
if (cur_sum > R) {
answer += cur->cheap_help(cur_sum - R);
cur_sum = R;
}
cur->trim(cur_sum - L);
}
return answer;
}
long long query(i32 L, i32 R) {
if (St45::active)
return St45::solve(L, R);
return solve(L, R);
}
