/**
* What an interesting pratice contest problem! I could only solve the first 3
* subtasks. (they're really trivial...) After that I searched online for an answer.
*
* The idea is like "Chain Reactions" (Google Code Jam 2022 Qual.). Since a node can
* only be chosen when its parent has been, we try to carefully merge leaf nodes into
* their parent through a greedy approach. Previously, I tried too hard on
* identifying the first / last selected node, which can be quite difficult when our
* search base is the whole tree.
*
* Each time, we try to find a node v with the smallest d[]/u[]. (This node need not
* be a leaf.) We try to merge v into its parent p, i.e. u[p] += u[v], d[p] += d[v],
* adjust the cost and let children of v have p as their direct parent. Why is
* this greedy correct?
* Although it could be difficult to determine which leaf is the last / which node
* is the first, if v has the smallest d[]/u[], there exists a optimal schedule in
* which v is built immediately after p is built. Simple argument of "exchange" can
* prove this. In such case, we can just consider the time diff. (and add the
* additional cost). Note that all children of v now have p as their direct parent
* because once v and p are built, all children of v can start.
* Learnt lots of stuff from this problem!
*
* PS1 The solution above is inspired by Tutis (from Codeforces) in
* https://codeforces.com/blog/entry/68632.
* PS2 When we apply our greedy strategy, the tree may be splited, forming a
* forest. We can add a pseudo root (the original -1 in parents[]) to simplify
* our code.
*
* Time Complexity: O(n * log(n))
* Implementation 2 (full solution, inspired by Tutis from Codeforces)
*/
#include <bits/stdc++.h>
#include "job.h"
typedef long long ll;
typedef std::vector<int> vec;
// DSU that can assign a value to each set
struct DSU {
vec parents, rank, values;
DSU(int n) {
parents.resize(n);
rank.resize(n);
values.resize(n);
for (int i = 0; i < n; i++)
parents[i] = i, rank[i] = 1;
}
inline int find(int k) {
if (parents[k] == k)
return k;
return parents[k] = find(parents[k]);
}
inline int& get_val(int k) {
k = find(k);
return values[k];
}
bool merge(int a, int b) { // merge b into a
a = find(a), b = find(b);
if (a == b)
return false;
int val = values[a];
if (rank[b] > rank[a])
std::swap(a, b);
rank[a] += rank[b], parents[b] = a, values[a] = val;
return true;
}
};
struct r_t {
int node;
double r;
};
inline bool operator<(const r_t& r1, const r_t& r2) { return r1.r > r2.r; }
ll scheduling_cost(std::vector<int> _p, std::vector<int> _u, std::vector<int> _d) {
int n = _p.size();
std::vector<int> parents(n + 1), u(n + 1), d(n + 1);
for (int i = 0; i < n; i++)
parents[i + 1] = _p[i] + 1, u[i + 1] = _u[i], d[i + 1] = _d[i];
parents[0] = -1, u[0] = d[0] = 0; // pseudo node
// Use dsu to maintain merged nodes
DSU dsu(n + 1);
dsu.get_val(0) = 0;
std::vector<double> r(n + 1);
std::priority_queue<r_t> pq;
for (int i = 1; i <= n; i++) {
dsu.get_val(i) = i, r[i] = double(d[i]) / u[i];
pq.push(r_t{i, r[i]});
}
ll cost = 0;
while (!pq.empty()) {
r_t rt = pq.top();
pq.pop();
if (dsu.get_val(rt.node) != rt.node || r[rt.node] != rt.r)
continue;
int v = rt.node, p = dsu.get_val(parents[v]);
// std::cerr << "(p,v): (" << p << ',' << v << "),\tpairs of (u,d) (" << u[p]
// << ',' << d[p] << ") (" << u[v] << ',' << d[v] << ")" << std::endl;
cost -= ll(d[v]) * ll(u[p]);
d[p] += d[v], u[p] += u[v], r[p] = double(d[p]) / u[p];
dsu.merge(p, v);
if (p > 0)
pq.push(r_t{p, r[p]});
}
cost += ll(u[0]) * ll(d[0]);
return cost;
}
