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 X first
#define Y second
#define ll long long
#define pb push_back
#define lch i << 1
#define rch i << 1 | 1
#define all(x) x.begin(),x.end()
const int N = 1e5 + 1;
typedef pair<int,int> ii;
namespace _IT {
int t[N << 2];
int add(int x,int y) {
return x == y ? x : 0;
}
int upd(int i,int l,int r,int L,int R,int v) {
if (l > R || L > r) return 0;
if (L <= l && r <= R) {
t[i] = v;
return 1;
}
int m = (l + r) / 2;
upd(lch,l,m,L,R,v); ++m;
upd(rch,m,r,L,R,v);
t[i] = add(t[lch],t[rch]);
}
int get(int i,int l,int r,int L,int R) {
if (t[i] && l < r)
t[lch] = t[i],
t[rch] = t[i];
if (l > R || L > r) return -1;
if (L <= l && r <= R) return t[i];
int m = (l + r) / 2;
int tmp1 = get(lch,l,m,L,R); ++m;
int tmp2 = get(rch,m,r,L,R);
if (tmp1 < 0) tmp1 = tmp2;
if (tmp2 < 0) tmp2 = tmp1;
return add(tmp1,tmp2);
}
int lef(int i,int l,int r,int p,int v) {
if (t[i] && l < r)
t[lch] = t[i],
t[rch] = t[i];
if (t[i]) return t[i] == v ? l : N;
int m = (l + r + 2) / 2;
if (m > p) return lef(lch,l,m - 1,p,v);
int x = lef(rch,m,r,p,v);
if (x == m--) {
x = lef(lch,l,m,m,v);
if (x == N)
x = m + 1;
}
return x;
}
};
namespace BIT {
vector<int> val;
vector<int> bit;
void add(int x) { val.pb(x); }
void ini() {
sort(all(val));
val.resize(unique(all(val)) - val.begin());
bit.assign(val.size() + 5,0);
}
void upd(int p,int v) {
p = upper_bound(all(val),p) - val.begin() + 1;
int K = bit.size();
for(; p < K ; p += p & -p)
bit[p] += v;
}
int get(int p) {
int ans = 0;
p = upper_bound(all(val),p) - val.begin();
for(; p > 0 ; p -= p & -p)
ans += bit[p];
return ans;
}
};
vector<int> g[N];
int nCh[N], led[N];
int pos[N], par[N];
int arr[N], tot = 0;
int dfs(int u,int p) {
par[u] = p;
nCh[u] = 1;
for(int v : g[u])
dfs(v,u),
nCh[u] += nCh[v];
return 1;
}
int hld(int u,int S) {
if (S) led[u] = u;
else led[u] = led[par[u]];
arr[++tot] = u;
pos[u] = tot;
int B = 0;
for(int v : g[u]) if (nCh[v] > nCh[B])
B = v;
if (B) hld(B,0);
for(int v : g[u]) if (v != B)
hld(v,1);
return 1;
}
int c[N];
int a[N], b[N];
int main() {
ios_base::sync_with_stdio(0);
cin.tie(0); cout.tie(0);
int n; cin >> n;
for(int i = 1 ; i <= n ; ++i) cin >> c[i];
for(int i = 2 ; i <= n ; ++i) {
cin >> a[i] >> b[i];
g[a[i]].pb(b[i]);
}
dfs(1,0);
hld(1,1);
for(int i = 1 ; i <= n ; ++i)
_IT::upd(1,1,n,pos[i],pos[i],c[i]);
for(int i = 2 ; i <= n ; ++i) {
vector<ii> vec;
BIT::val.clear();
for(int x = a[i] ; x ;) {
int r = pos[x];
int C = _IT::get(1,1,n,r,r);
int l = _IT::lef(1,1,n,r,C);
if (l < pos[led[x]])
l = pos[led[x]];
x = par[arr[l]];
_IT::upd(1,1,n,l,r,c[b[i]]);
BIT::add(C);
vec.pb({C,r - l + 1});
}
BIT::ini();
ll ans = 0;
for(ii p : vec) {
ans += 1ll * p.Y * BIT::get(p.X);
BIT::upd(p.X,p.Y);
}
cout << ans << "\n";
}
}
Compilation message (stderr)
construction.cpp: In function 'int _IT::upd(int, int, int, int, int, int)':
construction.cpp:33:5: warning: control reaches end of non-void function [-Wreturn-type]
}
^
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |