#include <bits/stdc++.h>
#define int long long
#define pii pair<int,int>
#define vi vector<int>
#define ff first
#define ss second
#define sp << " " <<
#define all(x) x.begin(),x.end()
#define big(x) ((int)(x.size()))
using namespace std;
const int MOD = 1e9+7, LIM = 1e6+1, inf = 2e9;
const int N = 1e5+1;
vi edges[N],sub(N),tp(N),par(N),dep(N);
struct FT {
vi bit;
int n;
vector<pii> upds;
FT(int nn) {
n = nn;
bit.assign(n + 1, 0);
}
void put(int p, int v,int r = 0) {
if (!r) upds.push_back({p,v});
for (int i = p; i <= n; i += i & -i) bit[i] += v;
}
int get(int p) {
int ans = 0;
for (int i = p; i > 0; i -= i & -i) ans += bit[i];
return ans;
}
int get(int l, int r) {
return get(r) - get(l - 1);
}
void reset() {
for (auto it : upds) put(it.ff,-it.ss,1);
upds.clear();
}
};
void dfs(int node,int p) {
if (node == p) dep[node] = 0;
else dep[node] = dep[p]+1;
par[node] = p;
sub[node] =1 ;
for (auto it : edges[node]) {
if (it == p) continue;
dfs(it,node);
sub[node]+=sub[it];
}
}
void hld(int node,int t) {
tp[node] = t;
int hvy = -1;
for (auto it : edges[node]) {
if (hvy == -1 || sub[it] > sub[hvy]) hvy = it;
}
if (hvy == -1) return;
hld(hvy,t);
for (auto it : edges[node]) if (it != hvy) hld(it,it);
}
vector<pii> stuf[N];
void solve() {
int n;
cin >> n;
vi c(n+1);
vi v;
for (int i=1;i<=n;i++) {
cin >> c[i];
v.push_back(c[i]);
}
sort(all(v));
v.erase(unique(all(v)),v.end());
for (int i=1;i<=n;i++) c[i] = upper_bound(all(v),c[i])-v.begin();
vector<pii> edg;
for (int i=1;i<n;i++) {
int a,b;
cin >> a >> b;
edg.push_back({a,b});
edges[a].push_back(b);
}
dfs(1,1);
hld(1,1);
vector<pii> chns;
FT ft(n);
auto qry = [&](int node,int v) -> int {
int cur = node;
int ans = 0;
ft.reset();
chns.clear();
while (1) {
int chain = tp[cur];
chns.push_back({chain,cur});
if (chain == 1) break;
cur = par[chain];
}
reverse(all(chns));
for (auto& [chain,curr] : chns) {
for (int i = big(stuf[chain])-1;i>=0;i--) {
int dp = dep[stuf[chain][i].ff];
int amt = (i == big(stuf[chain])-1)?(dp-dep[chain]+1):
(dp-dep[stuf[chain][i+1].ff]);
amt-=max(0ll,dp-dep[curr]);
amt = max(amt,0ll);
if (!amt) break;
ans+=amt*ft.get(stuf[chain][i].ss+1,n);
ft.put(stuf[chain][i].ss,amt);
}
}
return ans;
};
auto upd = [&](int node,int v) -> void {
int cur = node;
while (1) {
int chain = tp[cur];
while (big(stuf[chain]) && dep[stuf[chain].back().ff] <= dep[cur]) stuf[chain].pop_back();
stuf[chain].push_back({cur,v});
if (chain == 1) break;
cur = par[chain];
}
};
upd(1,c[1]);
for (int i=1;i<n;i++) {
auto[a,b] = edg[i-1];
cout << qry(a,c[a]) << '\n';
upd(b,c[b]);
}
}
signed main() {
ios_base::sync_with_stdio(0); cin.tie(0);
#ifdef Dodi
freopen("in.txt", "r", stdin);
freopen("out.txt", "w", stdout);
#endif
int t = 1;
//cin >> t;
while (t --> 0) solve();
}
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |