Submission #497589

#	Time	Username	Problem	Language	Result	Execution time	Memory
497589		Jarif_Rahman	Unique Cities (JOI19_ho_t5)	C++17	100 / 100	1750 ms	202544 KiB

joi2019_ho_t5

#include <bits/stdc++.h>
#define pb push_back
#define f first
#define sc second
using namespace std;
typedef long long int ll;
typedef string str;
const int k = 18;
const int inf = 1e8;
 
struct persistent_segtree{
    struct node{
        int sm;
        int l, r;
        node(){
            sm = 0LL;
            l = -1, r = -1;
        }
        node(int _sm){
            sm = _sm;
            l = -1, r = -1;
        }
        node(int _sm, int _l, int _r){
            l = _l, r = _r, sm = _sm;
        }
    };
    vector<node> v;
    int k;
    persistent_segtree(int n){
        k = 1;
        while(k <= n) k*=2; k*=2;
        v.resize(k, node(0));
        for(int i = 1; 2*i+1 < k; i++) v[i].l = 2*i, v[i].r = 2*i+1;
    }
    int upd(int nd, int a, int b, int in, int x){
        if(a == b){
            v.pb(node(x));
            return (int)v.size() - 1;
        }
        int md = (a+b)/2;
        if(in <= md){
            int rt = upd(v[nd].l, a, md, in, x);
            v.pb(node(v[v[nd].r].sm+v[rt].sm, rt, v[nd].r));
            return (int)v.size()-1;
        }
        else{
            int rt = upd(v[nd].r, md+1, b, in, x);
            v.pb(node(v[v[nd].l].sm+v[rt].sm, v[nd].l, rt));
            return (int)v.size()-1;
        }
    }
    int upd(int nd, int in, int x){
        return upd(nd, 0, k/2-1, in, x);
    }
    int sum(int l, int r, int nd, int a, int b){
        if(a > r || b < l){
            return 0;
        }
        if(a >= l && b <= r){
            return v[nd].sm;
        }
        int md = (a+b)/2;
        return sum(l, r, v[nd].l, a, md) + sum(l, r, v[nd].r, md+1, b);
    }
    int sum(int nd, int l, int r){
        return sum(l, r, nd, 0, k/2-1);
    }
};
 
int n, m;
vector<int> C;
vector<vector<int>> v;
vector<vector<pair<int, int>>> anc;
vector<int> cnt;
vector<int> pss;
vector<int> mxh;
vector<int> ans;
 
persistent_segtree ps(0);
 
pair<int, int> next(int nd, int h){
    if(h == 0) return {nd, 0};
    for(int i = k-1; i >= 0; i--){
        if(anc[nd][i].sc <= h){
            h-=anc[nd][i].sc;
            auto nxt = next(anc[nd][i].f, h);
            nxt.sc+=anc[nd][i].sc;
            return nxt;
        }
    }
    return {anc[nd][0].f, anc[nd][0].sc};
}
 
void dfs0(int nd, int ss){
    for(int x: v[nd]) if(x != ss) dfs0(x, nd);
    int mx = 0, smx = 0, in = n;
    for(int x: v[nd]) if(x != ss){
        if(mxh[x] > mx) smx = mx, mx = mxh[x], in = x;
        else if(mxh[x] > smx) smx = mxh[x];
    }
    mxh[nd] = mxh[in]+1;
    auto nxt = next(in, smx);
    nxt.sc++;
    nxt.sc = min(nxt.sc, inf);
    anc[nd][0] = nxt;
    for(int i = 1; i < k; i++){
        anc[nd][i].f = anc[anc[nd][i-1].f][i-1].f;
        anc[nd][i].sc = anc[nd][i-1].sc + anc[anc[nd][i-1].f][i-1].sc;
        anc[nd][i].sc = min(anc[nd][i].sc, inf);
    }
    cnt[nd] = cnt[nxt.f];
    pss[nd] = pss[nxt.f];
    if(ps.sum(pss[nd], C[nd], C[nd]) == 0){
        cnt[nd]++;
        pss[nd] = ps.upd(pss[nd], C[nd], 1);
    }
}
 
void dfs1(int nd, int ss){
    vector<pair<int, int>> sth;
    for(int x: v[nd]) sth.pb({mxh[x], x});
    sort(sth.rbegin(), sth.rend());
    while(sth.size() < 3) sth.pb({0, n});
 
    auto nxt = next(sth[0].sc, sth[1].f);
    ans[nd] = cnt[nxt.f];
 
    auto old_anc = anc[nd];
    int old_cnt = cnt[nd], old_mxh = mxh[nd];
    auto old_pss = pss[nd];
    for(int x: v[nd]) if(x != ss){
        pair<int, int> nxt;
        if(x == sth[0].sc) nxt = next(sth[1].sc, sth[2].f), mxh[nd] = mxh[sth[1].sc]+1;
        else if(x == sth[1].sc) nxt = next(sth[0].sc, sth[2].f), mxh[nd] = mxh[sth[0].sc]+1;
        else nxt = next(sth[0].sc, sth[1].f), mxh[nd] = mxh[sth[0].sc]+1;
 
        nxt.sc++;
        nxt.sc = min(nxt.sc, inf);
 
        anc[nd][0] = nxt;
        for(int i = 1; i < k; i++){
            anc[nd][i].f = anc[anc[nd][i-1].f][i-1].f;
            anc[nd][i].sc = anc[nd][i-1].sc + anc[anc[nd][i-1].f][i-1].sc;
            anc[nd][i].sc = min(anc[nd][i].sc, inf);
        }
        cnt[nd] = cnt[nxt.f];
        pss[nd] = pss[nxt.f];
        if(ps.sum(pss[nd], C[nd], C[nd]) == 0){
            cnt[nd]++;
            pss[nd] = ps.upd(pss[nd], C[nd], 1);
        }
        dfs1(x, nd);
    }
    anc[nd] = old_anc;
    cnt[nd] = old_cnt;
    mxh[nd] = old_mxh;
    pss[nd] = old_pss;
}
 
 
int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> n >> m;
    C.resize(n);
    v.resize(n);
    anc.resize(n+1, vector<pair<int, int>>(k, {n, inf}));
    cnt.resize(n+1, 0);
    mxh.resize(n+1, 1);
    mxh[n] = 0;
    ans.resize(n);
 
    for(int i = 0; i < n-1; i++){
        int a, b; cin >> a >> b; a--, b--;
        v[a].pb(b);
        v[b].pb(a);
    }
 
    for(int &x: C) cin >> x, x--;
 
    ps = persistent_segtree(m);
    pss.resize(n+1, 1);
 
    dfs0(0, -1);
    dfs1(0, -1);
 
    for(int x: ans) cout << x << "\n";
}

Compilation message (stderr)

joi2019_ho_t5.cpp: In constructor 'persistent_segtree::persistent_segtree(int)':
joi2019_ho_t5.cpp:31:9: warning: this 'while' clause does not guard... [-Wmisleading-indentation]
   31 |         while(k <= n) k*=2; k*=2;
      |         ^~~~~
joi2019_ho_t5.cpp:31:29: note: ...this statement, but the latter is misleadingly indented as if it were guarded by the 'while'
   31 |         while(k <= n) k*=2; k*=2;
      |                             ^

• Subtask 1: 4 / 4
• Subtask 2: 32 / 32
• Subtask 3: 32 / 32
• Subtask 4: 32 / 32