This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define pb push_back
#define pf push_front
#define sz(a) (int)a.size()
#define all(a) begin(a),end(a)
const int mxN = (int)2e5+2;
using ar = array<int,2>;
int n, tim=1;
vector<int> adj[mxN];
deque<ar> edge, ch[mxN];
int c[mxN], h[mxN], sub[mxN], p[mxN], st[mxN], en[mxN], dep[mxN], fen[mxN];
void upd(int x, int v){ for(; x<mxN; x+=x&-x) fen[x]+=v; }
int sum(int x){ int s=0; for(; x; x-=x&-x) s+=fen[x]; return s; }
int dfsSz(int s, int par){
sub[s] = 1; dep[s] = (par==-1?0:dep[par]+1); p[s] = par;
for(auto &u : adj[s]){
if(u==par) continue;
sub[s]+=dfsSz(u,s);
int x = adj[s][0];
if(x==par or sub[x]<sub[u]) swap(adj[s][0],u);
}
return sub[s];
}
void dfsHld(int s, int par, int head){
h[s] = head; st[s]=tim++;
for(auto u : adj[s]){
if(u==par) continue;
dfsHld(u,s,u==adj[s][0]?head:u);
}
en[s] = tim;
}
int calc(int a){
vector<ar> v; v.clear(); int ans = 0;
while(a!=-1){
int H = h[a], siz = dep[a]-dep[h[a]]+1;
vector<ar> tv; tv.clear();
for(auto [x,cnt] : ch[H])
if(siz) tv.pb({x,min(cnt,siz)}), siz-=min(cnt,siz);
reverse(all(tv)); for(auto u : tv) v.pb(u);
a = p[H];
}
for(int i = 1; i <= n; i++) fen[i]=0;
for(auto [x,y] : v) ans+=y*sum(x-1), upd(x,y);
return ans;
}
void addEdge(int b){
int v = c[b];
while(b!=-1){
int H = h[b], siz = dep[b]-dep[h[b]]+1; int tsz = siz;
while(!ch[H].empty() and siz){
int cnt = ch[H].front()[1];
if(cnt<=siz) ch[H].pop_front(), siz-=cnt;
else ch[H].front()[1]-=siz, siz=0;
}
ch[H].pf({v,tsz}); b = p[H];
}
}
int32_t main() {
cin >> n; vector<int> v;
for(int i = 1; i <= n; i++) cin >> c[i], v.pb(c[i]);
sort(all(v)), v.erase(unique(all(v)),end(v));
for(int i = 1; i <= n; i++) c[i] = lower_bound(all(v),c[i])-begin(v)+1;
for(int i = 1; i < n; i++){
int a, b; cin >> a >> b;
adj[a].pb(b), adj[b].pb(a), edge.pb({a,b});
}
ch[1].pb({c[1],1}); dfsSz(1,-1); dfsHld(1,-1,1);
for(auto [a,b] : edge) cout << calc(a) << "\n", addEdge(b);
}
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |