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>
#define MAXN 200007
using namespace std;
const int inf=1e9+7;
int n,c[MAXN],a[MAXN],b[MAXN],sz[MAXN],heavy[MAXN],parent[MAXN],cnt;
int pos[MAXN],head[MAXN],tt,lazy[4*MAXN],dep[MAXN];
pair<int,int> tree[4*MAXN],curr;
vector<int> v[MAXN],s;
vector< pair<int,int> > w;
map<int,int> mp;
void dfs(int x,int p,int d){
parent[x]=p; dep[x]=d;
sz[x]=1;
for(int i=0;i<v[x].size();i++){
if(v[x][i]==p)continue;
dfs(v[x][i],x,d+1);
if(sz[v[x][i]]>sz[heavy[x]])heavy[x]=v[x][i];
sz[x]+=sz[v[x][i]];
}
}
void decompose(int x,int p,int h){
tt++; pos[x]=tt; head[x]=h;
if(heavy[x]!=0){
decompose(heavy[x],x,h);
}
for(int i=0;i<v[x].size();i++){
if(v[x][i]==p or v[x][i]==heavy[x])continue;
decompose(v[x][i],x,v[x][i]);
}
}
void psh(int v){
if(lazy[v]==0)return;
tree[2*v].first=tree[2*v].second=lazy[v];
tree[2*v+1].first=tree[2*v+1].second=lazy[v];
lazy[2*v]=lazy[2*v+1]=lazy[v];
lazy[v]=0;
}
pair<int,int> combine(pair<int,int> fr,pair<int,int> sc){
return {min(fr.first,sc.first),max(fr.second,sc.second)};
}
void update(int v,int l,int r,int ll,int rr,int val){
if(ll>rr)return;
if(l==ll and r==rr){
tree[v].first=tree[v].second=val;
lazy[v]=val;
}else{
psh(v);
int tt=(l+r)/2;
update(2*v,l,tt,ll,min(tt,rr),val);
update(2*v+1,tt+1,r,max(tt+1,ll),rr,val);
tree[v]=combine(tree[2*v],tree[2*v+1]);
}
}
pair<int,int> info(int v,int l,int r,int ll,int rr){
if(ll>rr)return{inf,0};
if(l==ll and r==rr){
return tree[v];
}else{
psh(v);
int tt=(l+r)/2;
return combine( info(2*v,l,tt,ll,min(tt,rr)) , info(2*v+1,tt+1,r,max(tt+1,ll),rr) );
}
}
void upgrade(int x,int val){
while(x!=0){
update(1,1,n,pos[head[x]],pos[x],val);
x=parent[head[x]];
}
}
void add_segments(int l,int r){
while(r>=l){
int ll=l-1,rr=r,mid;
while(ll+1<rr){
mid=(ll+rr)/2;
curr=info(1,1,n,mid,r);
if(curr.first==curr.second){
rr=mid;
}else{
ll=mid;
}
}
curr=info(1,1,n,rr,r);
w.push_back({curr.first,r-rr+1});
r=rr-1;
}
}
int fenwick[MAXN];
void add(int pos,int val){
while(pos<=cnt){
fenwick[pos]+=val;
pos+=(pos & (-pos));
}
}
int getsum(int pos){
int res=0;
while(pos>=1){
res+=fenwick[pos];
pos-=(pos & (-pos));
}
return res;
}
long long query(int x){
long long ans=0;
w.clear();
while(x!=0){
add_segments(pos[head[x]],pos[x]);
x=parent[head[x]];
}
for(int i=0;i<w.size();i++){
ans+=(long long) getsum(w[i].first-1)*w[i].second;
add(w[i].first,w[i].second);
}
for(int i=0;i<w.size();i++){
add(w[i].first,-w[i].second);
}
return ans;
}
int main(){
ios_base::sync_with_stdio(0);
cin.tie(0);
cout.tie(0);
cin>>n;
for(int i=1;i<=n;i++){
cin>>c[i];
s.push_back(c[i]);
}
sort(s.begin(),s.end());
for(int i=0;i<s.size();i++){
if(i==0 or s[i]!=s[i-1])cnt++;
mp[s[i]]=cnt;
}
for(int i=1;i<=n;i++)c[i]=mp[c[i]];
for(int i=1;i<=n-1;i++){
cin>>a[i]>>b[i];
v[a[i]].push_back(b[i]);
v[b[i]].push_back(a[i]);
}
dfs(1,0,1);
decompose(1,0,1);
for(int i=1;i<=n;i++){
update(1,1,n,pos[i],pos[i],c[i]);
}
for(int i=1;i<=n-1;i++){
cout<<query(a[i])<<"\n";
upgrade(b[i],c[b[i]]);
}
return 0;
}
Compilation message (stderr)
construction.cpp: In function 'void dfs(int, int, int)':
construction.cpp:20:18: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
20 | for(int i=0;i<v[x].size();i++){
| ~^~~~~~~~~~~~
construction.cpp: In function 'void decompose(int, int, int)':
construction.cpp:36:18: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
36 | for(int i=0;i<v[x].size();i++){
| ~^~~~~~~~~~~~
construction.cpp: In function 'long long int query(int)':
construction.cpp:140:18: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
140 | for(int i=0;i<w.size();i++){
| ~^~~~~~~~~
construction.cpp:145:18: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
145 | for(int i=0;i<w.size();i++){
| ~^~~~~~~~~
construction.cpp: In function 'int main()':
construction.cpp:165:18: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
165 | for(int i=0;i<s.size();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... |