Submission #701664

#	Time	Username	Problem	Language	Result	Execution time	Memory
701664		urosk	Construction of Highway (JOI18_construction)	C++14	100 / 100	1136 ms	36456 KiB

construction

#define here cerr<<"===========================================\n"
#define dbg(x) cerr<<#x<<": "<<x<<endl;
#include "bits/stdc++.h"
//#include <ext/pb_ds/tree_policy.hpp>
//#include <ext/pb_ds/assoc_container.hpp>
#define ld double
#define ll long long
#define llinf 100000000000000000LL // 10^17
#define pb push_back
#define popb pop_back
#define fi first
#define sc second
#define endl '\n'
#define pll pair<ll,ll>
#define pld pair<ld,ld>
#define sz(a) (ll)(a.size())
#define all(a) a.begin(),a.end()
#define ceri(a,l,r) {cerr<<#a<<": ";for(ll i_ = l;i_<=r;i_++) cerr<<a[i_]<< " ";cerr<<endl;}
#define cer(a) {cerr<<#a<<": ";for(ll x_ : a) cerr<<x_<< " ";cerr<<endl;}

#define daj_mi_malo_vremena ios_base::sync_with_stdio(false);cerr.tie(0);cout.tie(0);cin.tie(0);

using namespace std;
//using namespace __gnu_pbds;
/*
ll add(ll x,ll y){
    x+=y;
    if(x<0){
        x%=mod;
        x+=mod;
    }else{
        if(x>=mod) x%=mod;
    }
    return x;
}
ll mul(ll a,ll b){
	ll ans = (a*b)%mod;
	if(ans<0) ans+=mod;
	return ans;
}
typedef tree<int,null_type,less<ll>,rb_tree_tag,tree_order_statistics_node_update> ordered_set;
typedef tree<int,null_type,less_equal<ll>,rb_tree_tag,tree_order_statistics_node_update> ordered_multiset;
mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());
ll rnd(ll l,ll r){
    return uniform_int_distribution<ll>(l,r)(rng);
}
*/
#define maxn 100005
ll n;
vector<ll> g[maxn];
vector<pll> v[maxn];
pll e[maxn];
ll c[maxn],sub[maxn],par[maxn],b[maxn],pos[maxn],siz[maxn],f[maxn];
ll dfssub(ll u,ll p){
    sub[u] = 1;
    for(ll s : g[u]) if(s!=p) sub[u]+=dfssub(s,u);
    return sub[u];
}
void dfs(ll u,ll p,ll up){
    v[up].pb({1,c[u]});
    ll c = 0;
    par[u] = p;
    b[u] = up;
    siz[up]++;
    pos[u] = sz(v[up]);
    for(ll s : g[u]) if(s!=p&&sub[s]>sub[c]) c = s;
    if(c==0) return;
    dfs(c,u,up);
    for(ll s : g[u]){
        if(s==p||s==c) continue;
        dfs(s,u,s);
    }
}
ll t[2*maxn],ls[2*maxn],rs[2*maxn],tsz = 0,root = 0,mx = 0;
void upd(ll v,ll tl,ll tr,ll i,ll x){
    if(tl==tr){t[v]+=x;return;}
    ll mid = (tl+tr)/2;
    if(i<=mid) upd(ls[v],tl,mid,i,x);
    else upd(rs[v],mid+1,tr,i,x);
    t[v] = t[ls[v]] + t[rs[v]];
}
ll sum(ll v,ll tl,ll tr,ll l,ll r){
    if(l>r||tl>tr||tl>r||tr<l) return 0;
    if(tl>=l&&tr<=r) return t[v];
    ll mid = (tl+tr)/2;
    return sum(ls[v],tl,mid,l,r) + sum(rs[v],mid+1,tr,l,r);
}
void init(ll &v,ll tl,ll tr){
    if(!v) v = ++tsz;
    if(tl==tr) return;
    ll mid = (tl+tr)/2;
    init(ls[v],tl,mid);
    init(rs[v],mid+1,tr);
}
void tc(){
    cin >> n;
    for(ll i = 1;i<=n;i++) cin >> c[i];
    {
        set<ll> s; map<ll,ll> mp; ll it = 0;
        for(ll i = 1;i<=n;i++) s.insert(c[i]);
        for(ll x : s) mp[x] = ++it;
        for(ll i = 1;i<=n;i++) c[i] = mp[c[i]];
        mx = it+1;
    }
    for(ll i = 1;i<=n-1;i++){
        ll x,y; cin >> x >> y;
        e[i] = {x,y};
        g[x].pb(y);
        g[y].pb(x);
    }
    dfssub(1,1);
    dfs(1,1,1);
    init(root,1,mx);
    stack<pll> st;
    for(ll i = 1;i<=n;i++) reverse(all(v[i]));
    for(ll i = 1;i<=n-1;i++){
        ll x = e[i].fi,y = e[i].sc;
        if(par[x]==y) swap(x,y);
        vector<ll> w;
        while(1){
            w.pb(b[x]);
            f[b[x]] = x;
            if(b[x]==1) break;
            x = par[b[x]];
        }
        reverse(all(w));
        x = e[i].fi^e[i].sc^y;
        ll ans = 0;
        for(ll z : w){
            ll cnt = 0;
            ll x = f[z];
            for(ll j = sz(v[z])-1;j>=0;j--){
                pll p = v[z][j];
                cnt+=p.fi;
                if(cnt<pos[x]){
                    ans+=p.fi*sum(root,1,mx,p.sc+1,mx);
                    upd(root,1,mx,p.sc,p.fi);
                    st.push(p);
                }else{
                    ans+=(p.fi-(cnt-pos[x]))*sum(root,1,mx,p.sc+1,mx);
                    upd(root,1,mx,p.sc,p.fi-(cnt-pos[x]));
                    st.push({p.fi-(cnt-pos[x]),p.sc});
                    break;
                }
            }
        }
        for(ll z : w){
            ll x = f[z];
            ll cnt = 0;
            for(ll j = sz(v[z])-1;j>=0;j--){
                pll p = v[z][j];
                cnt+=p.fi;
                v[z].popb();
                if(cnt>=pos[x]){
                    if(cnt!=pos[x]) v[z].pb({cnt-pos[x],p.sc});
                    break;
                }
            }
            v[z].pb({pos[x],c[y]});
        }
        while(sz(st)){
            pll p = st.top();
            st.pop();
            upd(root,1,mx,p.sc,-p.fi);
        }
        cout<<ans<<endl;
    }
}
int main(){
	daj_mi_malo_vremena
    int t; t = 1;
    while(t--){
        tc();
    }
	return 0;
}
/**
5
1 2 3 4 5
1 2
2 3
2 4
3 5

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
**/

• Subtask 1: 7 / 7
• Subtask 2: 9 / 9
• Subtask 3: 84 / 84