#include "bits/stdc++.h"
using namespace std;
#ifdef DEBUG
auto&operator<<(auto &o, pair<auto, auto> p) {o << "(" << p.first << ", " << p.second << ")"; return o;}
auto operator<<(auto&o,auto x)->decltype(x.end(),o){o<<"{"; for(auto e : x) o<<e<<", "; return o<<"}";}
#define debug(X) cerr << "["#X"]: " << X << '\n';
#else
#define cerr if(0)cout
#define debug(X) ;
#endif
using ll = long long;
#define all(v) (v).begin(), (v).end()
#define ssize(x) int(x.size())
#define fi first
#define se second
#define mp make_pair
#define eb emplace_back
struct segtree {
int base;
vector<int> t;
segtree() {}
void resize(int n) {
base = 1;
while (base < n) base *= 2;
t.resize(2 * base);
}
void update(int v, int x) {
v += base;
while (v)
t[v] += x, v /= 2;
}
int query(int l, int r) {
int res = 0;
l += base - 1; r += base + 1;
while (l/2 != r/2) {
if (l%2 == 0) res += t[l+1];
if (r%2 == 1) res += t[r-1];
l /= 2; r /= 2;
}
return res;
}
};
struct HLD {
int n;
vector<vector<int>> g;
vector<int> d, path, parent, sz, t;
vector<set<pair<pair<int, int>, int>>> s;
segtree T;
HLD(const vector<vector<int>> &graph): n(ssize(graph)), g(graph) {
d = path = parent = sz = t = vector<int>(n, -1);
T.resize(n);
s.resize(n);
d[0] = 0;
dfs(0, -1);
path[0] = 0;
make_hld(0, -1);
debug(parent);
debug(path);
debug(d);
}
void dfs(int v, int p) {
parent[v] = p;
sz[v] = 1;
t[v] = -1;
for (int u : g[v]) {
if (u == p) continue;
d[u] = d[v] + 1;
dfs(u, v);
sz[v] += sz[u];
t[v] = t[v] == -1 ? u : sz[t[v]] > sz[u] ? t[v] : u;
}
}
void make_hld(int v, int p) {
if (t[v] != -1) {
path[t[v]] = path[v];
make_hld(t[v], v);
}
for (int u : g[v]) {
if (u == p || u == t[v]) continue;
path[u] = u;
make_hld(u, v);
}
}
void update(int v, int c) {
while (v != -1) {
int p_start = path[v];
pair<pair<int, int>, int> last;
while (ssize(s[p_start]) && d[s[p_start].begin()->fi.fi] <= d[v]) {
last = *s[p_start].begin();
s[p_start].erase(s[p_start].begin());
}
if (d[last.fi.se] > d[v]) {
last.fi.fi = d[v] + 1;
s[p_start].insert(last);
}
s[p_start].insert({{d[p_start], d[v]}, c});
v = parent[p_start];
}
}
ll query(int v) {
debug(s);
vector<pair<int, int>> p;
while (v != -1) {
int p_start = path[v];
auto it = s[p_start].begin();
vector<pair<int, int>> p1;
while (it != s[p_start].end() && d[it->fi.fi] <= d[v]) {
p1.eb(min(d[v], it->fi.se) - it->fi.fi + 1, it->se);
++it;
}
for (int i = ssize(p1) - 1; i >= 0; --i) p.eb(p1[i]);
v = parent[p_start];
}
ll res = 0;
debug(p)
for (auto [cnt, x] : p) {
res += ll(cnt) * T.query(0, x-1);
T.update(x, cnt);
}
for (auto [cnt, x] : p)
T.update(x, -cnt);
return res;
}
};
int main() {
ios_base::sync_with_stdio(false); cin.tie(nullptr);
int n;
cin >> n;
vector<int> c(n);
map<int, int> remap;
for (int &x : c) {
cin >> x;
remap[x];
}
int cnt = 0;
for (auto &e : remap)
e.se = cnt++;
for (int &x : c) x = remap[x];
vector<pair<int, int>> e(n-1);
vector<vector<int>> g(n);
for (auto &[a, b] : e) {
cin >> a >> b;
--a; --b;
g[a].eb(b);
g[b].eb(a);
}
HLD hld(g);
hld.update(0, c[0]);
for (auto [a, b] : e) {
cout << hld.query(a) << '\n';
hld.update(b, c[b]);
}
return 0;
}
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |