Submission #867204

#	Time	Username	Problem	Language	Result	Execution time	Memory
867204		azberjibiou	Construction of Highway (JOI18_construction)	C++17	100 / 100	239 ms	84564 KiB

construction

#include <bits/stdc++.h>
#define all(v) v.begin(), v.end()
#define pb push_back
#define lb lower_bound
#define gibon ios::sync_with_stdio(false); cin.tie(0);
#define fi first
#define se second
#define pii pair<int, int>
#define pll pair<ll, ll>
typedef long long ll;
using namespace std;
const int mxN=100030;
const int mxM=500004;
const int MOD=998244353;
const ll INF=8e18;
typedef struct fenwick{
    int fen[mxN], n;
    void set_n(int m){n=m;}
    void init(){for(int i=0;i<=n+1;i++) fen[i]=0;}
    void upd(int idx, int val){
        while(idx<=n)   fen[idx]+=val, idx+=(idx&(-idx));
    }
    int solv(int idx){
        int res=0;
        while(idx)  res+=fen[idx], idx-=(idx&(-idx));
        return res;
    }
}fenwick;
int N;
int C[mxN];
vector <pii> E;
vector <int> v[mxN];
int dep[mxN], sub[mxN], g[mxN], par[mxN];
stack <pii> stk[mxN];
fenwick T1;
void input(){
    cin >> N;
    for(int i=1;i<=N;i++)   cin >> C[i];
    E.resize(N-1);
    for(auto &[x, y] : E)
    {
        cin >> x >> y;
        v[x].push_back(y);
        par[y]=x;
    }
}
void dfs1(int now){
    sub[now]=1;
    for(int nxt : v[now])
    {
        dep[nxt]=dep[now]+1;
        dfs1(nxt);
        sub[now]+=sub[nxt];
    }
}
void dfs2(int now){
    if(v[now].empty())  return;
    sort(all(v[now]), [](int a, int b){return sub[a]>sub[b];});
    for(int nxt : v[now])   g[nxt]=nxt;
    g[v[now][0]]=g[now];
    for(int nxt : v[now])   dfs2(nxt);
}
void heavy(int now, vector <pii> &ct, int col)
{
    vector <pii> nv;
    int idx=g[now], nl=dep[now]-dep[idx]+1;
    while(!stk[idx].empty())
    {
        auto [nc, len]=stk[idx].top();
        stk[idx].pop();
        if(nl>len)  nv.emplace_back(nc, len), nl-=len;
        else
        {
            nv.emplace_back(nc, nl);
            if(len!=nl) stk[idx].emplace(nc, len-nl);
            break;
        }
    }
    stk[idx].emplace(col, dep[now]-dep[idx]+1);
    reverse(all(nv));
    for(pii ele : nv)   ct.push_back(ele);
}
ll inv_count(vector <pii> &ct)
{
    ll res=0;
    vector <int> cr;
    cr.resize(ct.size());
    for(int i=0;i<ct.size();i++)    cr[i]=ct[i].fi;
    sort(all(cr));
    cr.erase(unique(all(cr)), cr.end());
    for(auto &[x, y] : ct)    x=lower_bound(all(cr), x)-cr.begin()+1;
    T1.set_n(cr.size()+1);
    T1.init();
    for(auto [col, cnt] : ct)
    {
        res+=(ll)T1.solv(col-1)*cnt;
        T1.upd(col, cnt);
    }
    return res;
}
ll qry(int now){
    vector <pii> ct;
    int np=now;
    if(g[now]==now) stk[now].emplace(C[now], 1), np=par[now];
    while(np){
        heavy(np, ct, C[now]);
        np=par[g[np]];
    }
    return inv_count(ct);
}
int main()
{
    gibon
    input();
    dfs1(1);
    g[1]=1;
    dfs2(1);
    stk[1].emplace(C[1], 1);
    for(auto [x, y] : E)    cout << qry(y) << '\n';
}

Compilation message (stderr)

construction.cpp: In function 'll inv_count(std::vector<std::pair<int, int> >&)':
construction.cpp:88: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]
   88 |     for(int i=0;i<ct.size();i++)    cr[i]=ct[i].fi;
      |                 ~^~~~~~~~~~

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