Submission #497091

#	Time	Username	Problem	Language	Result	Execution time	Memory
497091		Jarif_Rahman	Unique Cities (JOI19_ho_t5)	C++17	36 / 100	1213 ms	274436 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 = 17;
const int inf = 1e8;

struct persistent_segtree{
    struct node{
        int sm;
        node *l, *r;
        node(){
            sm = 0;
        }
        node(int x){
            sm = x;
        }
        node(node* _l, node* _r){
            l = _l, r = _r;
            sm = l->sm+r->sm;
        }
    };
    int k;
    persistent_segtree(int n){
        k = 1;
        while(k < n) k*=2;
        k*=2;
    }
    void init(node* nd, int kk){
        if(kk < 2) return;
        nd->l = new node();
        nd->r = new node();
        kk/=2;
        init(nd->l, kk);
        init(nd->r, kk);
    }
    node* init(){
        node* root = new node();
        init(root, k);
        return root;
    }
    node* upd(int in, node* nd, int a, int b, int x){
        if(a == b) return new node(x);
        int md = (a+b)/2;
        if(in <= md) return new node(upd(in, nd->l, a, md, x), nd->r);
        else return new node(nd->l, upd(in, nd->r, md+1, b, x));
    }
    node* upd(node* nd, int in, int x){
        return upd(in, nd, 0, k/2-1, x);
    }
    int sum(int l, int r, node* nd, int a, int b){
        if(a > r || b < l) return 0;
        if(a >= l && b <= r) return nd->sm;
        int md = (a+b)/2;
        return sum(l, r, nd->l, a, md)+sum(l, r, nd->r, md+1, b);
    }
    int sum(node* 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<persistent_segtree::node*> 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);
    auto init = ps.init();
    pss.resize(n+1, init);

    dfs0(0, -1);
    dfs1(0, -1);

    for(int x: ans) cout << x << "\n";
}

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