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;
long long a[100005], t[100005], sz[100005], depth[100005], par[100005], heavypar[100005], specialadj[100005], specialpar[100005], seg[400005], cnt;
vector<pair<long long, long long> > vec;
pair<long long, long long> edge[100005];
vector<long long> adj[100005];
void dfs1(long long u)
{
sz[u]=1;
for (long long v:adj[u])
{
depth[v]=depth[u]+1;
dfs1(v);
sz[u]+=sz[v];
}
}
void dfs2(long long u)
{
long long heavy=0;
t[u]=(++cnt);
for (long long v:adj[u])
if (!heavy || sz[heavy]<sz[v])
heavy=v;
if (heavy)
{
heavypar[heavy]=heavypar[u];
dfs2(heavy);
for (long long v:adj[u])
{
if (v==heavy)
continue;
heavypar[v]=v;
dfs2(v);
}
}
}
void push(long long i, long long l, long long r)
{
if (l==r || !seg[i])
return;
seg[i*2]=seg[i*2+1]=seg[i];
seg[i]=0;
}
void upd(long long i, long long l, long long r, long long ql, long long qr, long long x)
{
if (ql<=l && r<=qr)
{
seg[i]=x;
return;
}
push(i, l, r);
long long mid=(l+r)/2;
if (l<=qr && ql<=mid)
upd(i*2, l, mid, ql, qr, x);
if (mid<qr && ql<=r)
upd(i*2+1, mid+1, r, ql, qr, x);
}
long long qry(long long i, long long l, long long r, long long x)
{
if (l==r)
return seg[i];
push(i, l, r);
long long mid=(l+r)/2;
if (x<=mid)
return qry(i*2, l, mid, x);
else
return qry(i*2+1, mid+1, r, x);
}
long long calc(long long l, long long r)
{
if (l==r)
return 0;
long long mid=(l+r)/2, ind=0, sum=0, ans=0;
vector<pair<long long, long long> > vec1, vec2;
for (long long i=l; i<=mid; i++)
vec1.push_back(vec[i]);
for (long long i=mid+1; i<=r; i++)
vec2.push_back(vec[i]);
sort(vec1.begin(), vec1.end());
sort(vec2.begin(), vec2.end());
for (long long i=0; i<vec2.size(); i++)
{
while (ind<vec1.size() && vec1[ind].first<vec2[i].first)
{
sum+=vec1[ind].second;
ind++;
}
ans+=sum*vec2[i].second;
}
return ans+calc(l, mid)+calc(mid+1, r);
}
int main()
{
ios::sync_with_stdio(0);
cin.tie(0);
cout.tie(0);
long long n;
cin >> n;
for (long long i=1; i<=n; i++)
cin >> a[i];
for (long long i=1; i<n; i++)
{
long long u, v;
cin >> u >> v;
edge[i]={u, v};
adj[u].push_back(v);
par[v]=u;
}
par[1]=heavypar[1]=1;
dfs1(1);
dfs2(1);
for (long long i=1; i<=n; i++)
{
upd(1, 1, n, t[i], t[i], i);
specialpar[i]=i;
}
for (long long i=1; i<n; i++)
{
long long u=edge[i].first, specialgrp=qry(1, 1, n, t[u]);
vec.clear();
vec.push_back({a[specialgrp], depth[u]-depth[specialpar[specialgrp]]+1});
u=specialpar[specialgrp];
while (u!=1)
{
long long v=par[u], w=specialadj[v], specialgrp2=qry(1, 1, n, t[v]);
vec.push_back({a[specialgrp2], depth[u]-depth[specialpar[specialgrp2]]});
specialadj[v]=u;
u=specialpar[specialgrp]=specialpar[specialgrp2];
specialpar[specialgrp2]=w;
}
cout << calc(0, vec.size()-1) << '\n';
u=edge[i].first;
long long v=heavypar[u];
upd(1, 1, n, t[v], t[u], edge[i].second);
while (v!=1)
{
u=par[v];
v=heavypar[u];
upd(1, 1, n, t[v], t[u], edge[i].second);
}
u=edge[i].second, v=edge[i].first;
if (specialadj[v])
{
long long w=specialadj[v], specialgrp2=qry(1, 1, n, t[w]);
specialpar[specialgrp2]=w;
}
specialadj[v]=u;
specialpar[u]=1;
}
}
Compilation message (stderr)
construction.cpp: In function 'long long int calc(long long int, long long int)':
construction.cpp:81:26: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
81 | for (long long i=0; i<vec2.size(); i++)
| ~^~~~~~~~~~~~
construction.cpp:83:19: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
83 | while (ind<vec1.size() && vec1[ind].first<vec2[i].first)
| ~~~^~~~~~~~~~~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |