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 pii pair<int,int>
#define F first
#define S second
#define mp make_pair
#define pb push_back
const int N = 200005;
int sz[N] , in[N] , out[N] , C[N], t = 0 , ch[N] , depth[N] , inv[N] , pai[N] , endch[N];
vector<int> adj[N];
int n;
vector<pii> edge;
vector<pii> cor;
void dfs(int x , int y){
sz[x] = 1;
depth[x] = 0;
depth[x] = depth[y] + 1;
pai[x] = -1;
if(x != y) pai[x] = y;
for(auto &p : adj[x]){
if(p == y) continue;
dfs(p,x);
sz[x] += sz[p];
if(adj[x][0] == y || sz[p] >= sz[adj[x][0]]) swap(p, adj[x][0]);
}
}
int dfs_hld(int x, int y){
in[x] = ++t;
inv[in[x]] = x;
endch[ch[x]] = x;
for(auto p : adj[x]){
if(p == y) continue;
if(p == adj[x][0]) ch[p] = ch[x];
else ch[p] = p;
dfs_hld(p , x);
}
out[x] = t;
}
map<int,int> ordi;
set<int> sorti;
int BIT[N + 100];
void add(int x , int v){
for(int i = x ; i < (N + 100) ; i += (i&-i)) BIT[i] += v;
}
int get(int x){
int s = 0;
for(int i = x ; i > 0 ; i-= (i&-i)){
s += BIT[i];
}
return s;
}
deque< array<int, 4> > cs[N + 100];
int32_t main(){
scanf("%lld" , &n);
for(int i = 1 ; i <= n; i ++) scanf("%lld" , &C[i]) , sorti.insert(C[i]);
int CC = 0;
for(auto p : sorti){
ordi[p] = ++CC;
}
for(int i = 1 ; i <= n; i ++){
C[i] = ordi[C[i]];
}
for(int i = 1 ; i < n ;i ++){
int x , y;
scanf("%lld%lld" , &x , &y);
adj[x].pb(y);
adj[y].pb(x);
edge.push_back(pii(x,y));
}
ch[1] = 1;
dfs(1,1);
dfs_hld(1,1);
ch[1] = 1 , pai[1] = -1;
array<int, 4> ux;
ux[0] = in[1] , ux[1] = in[1] , ux[2] = C[1] , ux[3] = 1;
cs[ch[1]].push_back(ux);
for(auto p : edge){
vector<pii> q;
int X = p.F;
while(X != -1){
int l = 0 , r = cs[ch[X]].size();
r--;
int ansj = -1;
while(l<=r){
int mid = (l+r)/2;
if(cs[ch[X]][mid][0] <= in[X]){
ansj = mid;
l = mid + 1;
}
else
r = mid - 1;
}
if(ansj != -1){
q.push_back(pii(cs[ch[X]][ansj][1+1], in[X] - cs[ch[X]][ansj][0] + 1));
}
for(int i = ansj - 1 ; i >= 0 ; i--){
q.push_back(pii(cs[ch[X]][i][1+1] , cs[ch[X]][i][2+1]));
}
X = pai[ch[X]];
}
reverse(q.begin() , q.end());
long long ans = 0;
for(auto w : q){
add(w.F , w.S);
int U = get(N + 99) - get(w.F);
long long P = U;
P *= w.S;
ans += P;
}
for(auto w : q){
add(w.F , -w.S);
}
X = p.S;
while(X != -1){
array<int, 4> ux;
ux[0] = in[ch[X]] , ux[1] = in[X] , ux[2] = C[p.S] , ux[3] = 0;
if(X == p.S) ux[3]++;
while(cs[ch[X]].size() && cs[ch[X]].front()[1] <= in[X]) ux[3] += cs[ch[X]].front()[3] , cs[ch[X]].pop_front();
if(cs[ch[X]].size() && cs[ch[X]].front()[0] <= in[X]){
int R = in[X] - cs[ch[X]].front()[0] + 1;
cs[ch[X]].front()[3] -= R;
cs[ch[X]].front()[0] = in[adj[X][0]];
ux[3] += R;
}
while(cs[ch[X]].size() && cs[ch[X]].front()[2] == C[p.S]) ux[3] += cs[ch[X]].front()[3] , cs[ch[X]].pop_front();
cs[ch[X]].push_front(ux);
X = pai[ch[X]];
}
printf("%lld\n" , ans);
}
}
/*
10
1 7 3 4 8 6 2 9 10 5
1 2
1 3
2 4
3 5
2 6
3 7
4 8
5 9
6 10
*/
Compilation message (stderr)
construction.cpp: In function 'long long int dfs_hld(long long int, long long int)':
construction.cpp:39:1: warning: no return statement in function returning non-void [-Wreturn-type]
}
^
construction.cpp: In function 'int32_t main()':
construction.cpp:60:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
scanf("%lld" , &n);
~~~~~^~~~~~~~~~~~~
construction.cpp:61:54: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
for(int i = 1 ; i <= n; i ++) scanf("%lld" , &C[i]) , sorti.insert(C[i]);
~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~
construction.cpp:71:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
scanf("%lld%lld" , &x , &y);
~~~~~^~~~~~~~~~~~~~~~~~~~~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |