#include <bits/stdc++.h>
using namespace std;
#define ll long long
const ll mmod = 998244353;
#define vl vector<long long>
#define vll vector<vector<long long>>
#define pl pair<long long, long long>
#define vb vector<bool>
vll tree;
vl weights;
vl c;
vl parents;
vl w;
vl subs;
// DFS už nepotřebujeme, nahradí ho set
// pl dfs(ll v, ll p, ll l){ ... }
long long query(int L, int R){
ll n = tree.size();
c.clear(); c.resize(n, 0);
subs.clear(); subs.resize(n, 0);
// Pole setů: každý index nese set kandidátů {váha, id_vrcholu} pro daný podstrom
vector<set<pair<ll, ll>>> sets(n);
// Inicializace listů
for (ll i = 1; i < n; i++){
if (tree[i].size() == 1){
c[i] = L;
subs[i] = L;
}
}
// Bottom-up průchod
for (ll i = n - 1; i >= 0; i--){
ll suma = subs[i]; // Začneme s tím, co tam je (pro listy L, jinak 0)
// Přidáme aktuální vrchol do svého setu kandidátů
sets[i].insert({w[i], i});
for (ll j = 0; j < tree[i].size(); j++){
ll child = tree[i][j];
if (child != parents[i]){
suma += subs[child]; // OPRAVA: přičítáme subs dítěte, ne j-tého prvku pole subs
// "Hloupé" sloučení setů: prostě vše z dítěte nasypeme k rodiči
for (auto p : sets[child]){
sets[i].insert(p);
}
}
}
subs[i] = suma;
// Dokud je součet podstromu moc velký
while (subs[i] > R){
if (sets[i].empty()) break; // Pojistka, nemělo by nastat pokud existuje řešení
// Podíváme se na nejlevnější vrchol v podstromu
pair<ll, ll> best = *sets[i].begin();
ll pos = best.second;
// VALIDACE: Zkontrolujeme cestu nahoru k "i", zda můžeme odebírat
bool is_valid = true;
ll curr = pos;
while (true) {
if (subs[curr] <= L) {
is_valid = false;
break;
}
if (curr == i) break;
curr = parents[curr];
}
// Pokud vrchol není validní (někde na cestě bychom podtekli L), vyhodíme ho
if (!is_valid) {
sets[i].erase(sets[i].begin());
continue;
}
// Pokud validní je, zjistíme, o kolik můžeme snížit (k)
ll k = subs[i] - R; // Maximálně o tolik, kolik přečuhujeme přes R
curr = pos;
while (true){
k = min(k, subs[curr] - L); // Omezíme tím, kolik chybí do L na cestě
if (curr == i) break;
curr = parents[curr];
}
// Aplikujeme změnu
c[pos] -= k;
// Aktualizujeme součty na cestě nahoru k "i"
curr = pos;
while (true){
subs[curr] -= k;
if (curr == i) break;
curr = parents[curr];
}
// Poznámka: Vrchol ze setu nevyhazujeme hned, protože možná půjde
// v dalším kole snížit znovu. Vyhodí ho až validační check,
// až narazí na hranu L.
}
}
ll price = 0;
for (ll i = 0; i < n; i++){
price += w[i] * abs(c[i]);
}
return price;
}
void init(std::vector<int> P, std::vector<int> W){
// Reset globálních proměnných pro případ vícenásobného volání testeru
weights.clear();
tree.clear();
parents.clear();
w.clear();
for (ll i = 0; i < W.size(); i++){ weights.push_back(W[i]); }
tree.resize(P.size());
for (ll i = 1; i < P.size(); i++){
tree[i].push_back(P[i]);
tree[P[i]].push_back(i);
}
for (ll i = 0; i < P.size(); i++){
parents.push_back(P[i]);
w.push_back(W[i]);
}
}
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |