Submission #481123

#	Time	Username	Problem	Language	Result	Execution time	Memory
481123		K4YAN	Construction of Highway (JOI18_construction)	C++17	0 / 100	2 ms	3020 KiB

construction

//Be Name Khoda

#include<bits/stdc++.h>
#pragma GCC optimize ("Ofast,unroll-loops")

using namespace std;

#define ll long long
#define pb push_back
#define pii pair<int, int>
#define piii pair<pii, int>
#define piiii pair<pii, pii>
#define all(x) x.begin(), x.end()

const int mxn = 1e5 + 16;
int n, q, w;
int subt[mxn], h[mxn], hch[mxn], chn[mxn], lbl[mxn], par[mxn][17];
vector<int> g, adj[mxn];
vector<pii> queries;
bitset<mxn> mark;

struct item {

    ll ans, lme, rme, nlme, nrme, sz;
    //  0 ,  k ,  k , rx-lx, rx-lx, rx-lx;

};

struct segtree {
    int sz = 1;
    item NE = {0, 0, 0, 0, -1};
    vector<item> v;
    vector<int> lz;
    void init(int n) {
        while(sz < n) sz <<= 1;
        v.assign(2 * sz, NE);
        lz.assign(2 * sz, 0);
    }
    void shift(int x, int lx, int rx) {
        if(lz[x] == 0) return;
        v[x] = {0, lz[x], lz[x], rx-lx, rx-lx, rx-lx};
        if(rx - lx > 1) {
            lz[2 * x + 1] = lz[2 * x + 2] = lz[x];
        }
        lz[x] = 0; return;
    }
    item upd(item a, item b) {
        if(a.sz == -1) return b;
        if(b.sz == -1) return a;
        ll ans = a.ans + b.ans, s = a.sz + b.sz, x = a.nrme + b.nlme;
        if(a.rme == b.lme) {
            if(b.nlme == b.sz && a.nrme == a.sz) {
                return {0, a.lme, b.rme, s, s, s};
            }
            if(b.nlme != b.sz && a.nrme == a.sz) {
                return {ans, a.lme, b.rme, (a.sz + b.nlme), b.nrme, s};
            }
            if(b.nlme == b.sz && a.nrme != a.sz) {
                return {ans, a.lme, b.rme, a.nlme, (b.sz + a.nrme), s};
            }
            return {ans + (x * (x - 1) / 2), a.lme, b.rme, a.nlme, b.nrme, s};
        }
        else {
            if(b.nlme == b.sz && a.nrme == a.sz) {
                return {0, a.lme, b.rme, a.nlme, b.nrme, s};
            }
            if(b.nlme != b.sz && a.nrme == a.sz) {
                return {ans + (b.nlme * (b.nlme - 1) / 2), a.lme, b.rme, a.nlme, b.nrme, s};
            }
            if(b.nlme == b.sz && a.nrme != a.sz) {
                return {ans + (a.nrme * (a.nrme - 1) / 2), a.lme, b.rme, a.nlme, b.nrme, s};
            }
            return {ans + (a.nrme * (a.nrme - 1) / 2) + (b.nlme * (b.nlme - 1) / 2), a.lme, b.rme, a.nlme, b.nrme, s};
        }
    }
    void set(int l, int r, ll k, int x, int lx, int rx) {
        if(rx <= l || r <= lx) return;
        if(l <= lx && rx <= r) {
            lz[x] = k;
            shift(x, lx, rx);
            return;
        }
        shift(x, lx, rx);
        int m = (lx + rx) >> 1;
        set(l, r, k, 2 * x + 1, lx, m);
        set(l, r, k, 2 * x + 2, m, rx);
        v[x] = upd(v[2 * x + 1], v[2 * x + 2]);
    }
    void set(int l, int r, ll k) {
        set(l, r, k, 0, 0, sz);
        return;
    }
    item cal(int l, int r, int x, int lx, int rx) {
        if(r <= lx || rx <= l) return NE;
        shift(x, lx, rx);
        if(l <= lx && rx <= r) return v[x];
        int m = (lx + rx) >> 1;
        item ans = cal(l, r, 2 * x + 1, lx, m);
        item ans2 = cal(l, r, 2 * x + 2, m, rx);
        return ans = upd(ans, ans2);
    }
    item cal(int l, int r) {
        return cal(l, r, 0, 0, sz);
    }
};
segtree st;
item hp;

void DFS(int v) {

    mark.set(v), h[v] = w++, subt[v] = 1;
    for(auto u : adj[v]) {
        if(!mark[u]) {
            par[u][0] = v;
            DFS(u);
            subt[v] += subt[u];
            if(subt[hch[v]] < subt[u]) {
                hch[v] = u;
            }
        }
    }w--;

}

void SAGZAN(int v) {

    mark.set(v), lbl[v] = w++;
    if(chn[v] == -1) chn[v] = v;
    if(hch[v] != 0) {
        chn[hch[v]] = chn[v];
        SAGZAN(hch[v]);
    }
    for(auto u : adj[v]) {
        if(!mark[u]) {
            SAGZAN(u);
        } 
    }

}

void HLD_set(int v, int k) {
    int t, r = v;
    while(r > 0) {
        t = chn[r];
        if(t == 1) {
            st.set(0, lbl[r] + 1, k);
            break;
        }
        st.set(lbl[chn[r]], lbl[r] + 1, k);
        r = par[chn[r]][0];
    }
    return;
}

ll HLD_cal(int v) {

    ll t, r = v, cnt = 0, num = 0, ans = 0;
    while(r > 0) {
        t = chn[r];
        if(h[t] <= 0) {
            hp = st.cal(0, lbl[r] + 1);
            ans += hp.ans;
            if(hp.rme != num) {
                ans += (1LL * cnt * (cnt - 1) / 2);
                ans += (1LL * hp.nrme * (hp.nrme - 1) / 2);
                if(hp.nrme != hp.sz) ans += (1LL * hp.nlme * (hp.nlme - 1) / 2);
            }
            else {
                cnt += hp.nrme;
                ans += (1LL * cnt * (cnt - 1) / 2);
                if(hp.nrme != hp.sz) ans += (1LL * hp.nlme * (hp.nlme - 1) / 2);
            }
            break;
        }
        hp = st.cal(lbl[chn[r]], lbl[r] + 1);
        r = par[chn[r]][0];
        ans += hp.ans;
        if(hp.rme != num) {
            ans += (1LL * cnt * (cnt - 1) / 2);
            if(hp.nrme != hp.sz) ans += (1LL * hp.nrme * (hp.nrme - 1) / 2);
            cnt = hp.nlme, num = hp.lme;
        }
        else {
            cnt += hp.nrme;
            if(hp.nrme != hp.sz) ans += (1LL * cnt * (cnt - 1) / 2), cnt = hp.nlme;
            num = hp.lme;
        }
    }
    return ans;

}

void input() {
    
    cin >> n;
    for(int i = 0; i < n; i++) {
        cin >> q;
        g.push_back(q);
    }
    for(int i = 1; i < n; i++) {
        int v, u;
        cin >> v >> u;
        adj[v].push_back(u), adj[u].push_back(v);
        queries.push_back({v, u});
    }

}

void solve() {

    DFS(1), mark.reset(), fill(chn, chn + mxn, -1), SAGZAN(1);
    for(int i = 1; i < 17; i++) {
        for(int j = 1; j <= n; j++) {
            par[j][i] = par[par[j][i - 1]][i - 1];
        }
    }
    st.init(n);
    for(int i = 0; i < n - 1; i++) {
        int v = queries[i].first, u = queries[i].second;
        ll ans = 1LL * h[u] * (h[u] - 1) / 2, mtm;
        mtm = HLD_cal(v), ans -= mtm;
        HLD_set(u, g[u - 1]);
        g[0] = g[u - 1];
        cout << ans + 1 << "\n";
    }

}

int main() {

    ios_base::sync_with_stdio(false), cin.tie(NULL), cout.tie(NULL);

    input(), solve();

    return 0;

}
/*
5
1 2 3 4 5
1 2
2 3
2 4
3 5
*/

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