Submission #1316615

#TimeUsernameProblemLanguageResultExecution timeMemory
1316615olizarowskibConstruction of Highway (JOI18_construction)C++20
100 / 100
402 ms26724 KiB
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define int ll
#define ft first
#define sc second
#define pb push_back
#define all(x) x.begin(), x.end()
#define pi pair<int, int>
#define vi vector<int>
#define tii tuple<int, int, int>
#define tiii tuple<int, int, int, int>
#define tiiii tuple<int, int, int, int, int>
#define vpi vector<pi>
#define vtii vector<tii>
#define vtiii vector<tiii>

const int N = (1 << 20), MOD = 1e9 + 7, INF = int(1e18);

int p[N], w[N], d[N];
vi graf[N];

struct treep{
    int t[2 * N];
    vi z;

    void upd(int x, int v){
        z.pb(x);
        x += N;
        t[x] += v;
        do{
            x /= 2;
            t[x] = t[2 * x] + t[2 * x + 1];
        }while(x > 1);
    }

    int que(int l, int r){
        if(r < l) return 0;
        l += N - 1; r += N + 1;
        int ans = 0;
        while(l / 2 != r / 2){
            if(l % 2 == 0) ans += t[l + 1];
            if(r % 2 == 1) ans += t[r - 1];
            l /= 2; r /= 2;
        } 
        return ans;
    }

    void clr(){
        while(!z.empty()){
            int x = que(z.back(), z.back());
            upd(z.back(), -x);
            z.pop_back();
            z.pop_back();
        }
    }
};


struct trees{
    pi t[2 * N];
    int cz = 0;
    void upd(int l, int r, int x){
        cz++;
        l += N - 1; r += N + 1;
        while(l / 2 != r / 2){
            if(l % 2 == 0) t[l + 1] = {x, cz};
            if(r % 2 == 1) t[r - 1] = {x, cz};
            l /= 2; r /= 2;
        }
    }

    int que(int x){
        x += N;
        pi ans = {-1, -1};
        while(x >= 1){
            if(t[x].sc > ans.sc) ans = t[x];
            x /= 2;
        }
        return ans.ft;
    }
};

treep tri;
trees hld;


int idx[N], sz[N];

void dfs(int u, int v){
    p[u] = v;
    d[u] = d[v] + 1;
    sz[u] = 1;

    for(auto i : graf[u]){
        if(i == v) continue;
        dfs(i, u);
        sz[u] += sz[i];
    }
}

int hr[N];

void dfs2(int u, int v, int &t){
    idx[u] = ++t;
    int mx = -1;
    for(auto i : graf[u]){
        if(i == v) continue;
        if(sz[i] > sz[mx]) mx = i;
    }
    if(mx == -1) return;
    hr[mx] = hr[u];
    dfs2(mx, u, t);
    for(auto i : graf[u]){
        if(i == v || i == mx) continue;
        hr[i] = i;
        dfs2(i, u, t);
    }
}

int f[N], l[N];


int get_col(int x){
    return hld.que(idx[x]);
}

void upd_col(int x){
    int c = x;
    hld.upd(idx[x], idx[x], x);
    while(x > 1){
        int a = hr[x];
        if(a == p[x]) hld.upd(idx[a], idx[a], c);
        else hld.upd(idx[a], idx[x], c);
        x = a;
    }
}


int add(int x){
    int org = x;
    int akt = get_col(p[x]), lst = p[x], c = akt;
    int ans = 0;
    do{
        akt = f[akt];
        f[c] = l[lst];
        if(lst != x) l[lst] = x;
        int dl = d[lst] - d[akt] + 1;
        ans += dl * tri.que(1, w[c] - 1);
        tri.upd(w[c], dl);
        lst = p[akt];
        x = akt;
        c = get_col(lst);
        akt = c;
    }while(x != 1);
    upd_col(org);
    tri.clr();
    f[org] = 1;
    l[p[org]] = org;
    return ans;
}

signed main(){
    ios_base::sync_with_stdio(0);
    cin.tie(NULL); cout.tie(NULL);
    // random_device rd;
    // mt19937_64 gen(rd());
    int n;
    cin >> n;
    vpi v;
    for(int i = 1; i <= n; i++){
        int a;
        cin >> a;
        v.pb({a, i});
    }
    sort(all(v));
    for(int i = 0; i < n; i++){
        auto [a, b] = v[i];
        if(i == 0 || a != v[i - 1].ft) w[b] = i + 1;
        else w[b] = w[v[i - 1].sc];
    }
    vpi z;
    for(int i = 1; i < n; i++){
        int a, b;
        cin >> a >> b;
        z.pb({a, b});
        graf[a].pb(b); graf[b].pb(a);
    }
    f[1] = hr[1] = 1;
    int t = 0;
    dfs(1, 1);
    dfs2(1, 1, t);
    hld.upd(1, 1, 1);
    for(int i = 1; i <= n; i++){
        if(hr[i] == i) hr[i] = p[i];
    }
    for(auto i : z) cout << add(i.sc) << "\n";
    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...