제출 #1345843

#제출 시각아이디문제언어결과실행 시간메모리
1345843anteknneConstruction of Highway (JOI18_construction)C++20
100 / 100
615 ms84292 KiB
#include<bits/stdc++.h>
#include<ext/pb_ds/assoc_container.hpp>
#include<ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

typedef long long ll;
typedef long double ld;

#define pb push_back
#define pii pair<int, int>
#define pll pair<ll, ll>
#define st first
#define nd second
#define ordered_set tree<pii, null_type, less<pii>, rb_tree_tag, tree_order_statistics_node_update>
#define debug false

const int MAXN = 100 * 1000 + 17;
int ojc[MAXN];
int c[MAXN];
vector<int> graf[MAXN];
vector<int> vec;
int roz[MAXN];
int gora[MAXN];
int st[MAXN * 4];
vector<pii> kraw;
deque<pii> kol[MAXN];
int lev[MAXN];
vector<pii> operac;

void aktualizuj(int p, int k, int ind, int w, int i) {
    if (p == k && p == ind) {
        st[i] += w;
        return;
    }
    int sr = (p + k) / 2;
    if (ind <= sr) {
        aktualizuj(p, sr, ind, w, i * 2);
    } else {
        aktualizuj(sr + 1, k, ind, w, i * 2 + 1);
    }
    st[i] = (st[i * 2] + st[i * 2 + 1]);
}

int zapytanie(int p, int k, int a, int b, int i) {
    if (p > b || k < a) {
        return 0;
    }
    if (a <= p && k <= b) {
        return st[i];
    }
    int sr = (p + k) / 2;
    return (zapytanie(p, sr, a, b, i * 2) + zapytanie(sr + 1, k, a, b, i * 2 + 1));
}

void DFS (int v, int o) {
    roz[v] = 1;
    for (auto x : graf[v]) {
        if (x != o) {
            DFS(x, v);
            roz[v] += roz[x];
        }
    }
}

void DFS2 (int v, int o) {
    int u = -1, ile = 0;
    for (auto x : graf[v]) {
        if (x != o) {
            if (roz[x] > ile) {
                ile = roz[x];
                u = x;
            }
        }
    }
    if (u != -1) {
        gora[u] = gora[v];
    }
    for (auto x : graf[v]) {
        if (x != o && x != u) {
            gora[x] = x;
        }
    }

    for (auto x : graf[v]) {
        if (x != o) {
            DFS2(x, v);
        }
    }
}

int main () {
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);

    int n;
    cin >> n;

    for (int i = 1; i <= n; i ++) {
        cin >> c[i];
    }

    int a, b;
    for (int i = 0; i < n - 1; i ++) {
        cin >> a >> b;
        ojc[b] = a;
        lev[b] = lev[a] + 1;
        graf[a].pb(b);
        graf[b].pb(a);
        kraw.pb({a, b});
    }

    for (int i = 1; i <= n; i ++) {
        vec.pb(c[i]);
    }

    sort(vec.begin(), vec.end());
    vec.erase(unique(vec.begin(), vec.end()), vec.end());

    for (int i = 1; i <= n; i ++) {
        c[i] = lower_bound(vec.begin(), vec.end(), c[i]) - vec.begin();
    }

    int m = int(vec.size());

    gora[1] = 1;
    DFS(1, 1);
    DFS2(1, 1);

    kol[1].pb({c[1], 1});

    for (int i = 0; i < n - 1; i ++) {
        a = kraw[i].st;
        b = kraw[i].nd;

        ll wyn = 0;

        int v = a;
        operac.clear();
        while (true) {
            int g = gora[v];
            int lg = lev[g];
            vector<pii> jakie = {};
            while (!kol[g].empty() && (lg + kol[g].front().nd - 1) <= lev[v]) {
                jakie.pb(kol[g].front());
                operac.pb({kol[g].front().st, kol[g].front().nd});
                lg += kol[g].front().nd;
                kol[g].pop_front();
            }

            if (!kol[g].empty() && lg <= lev[v]) {
                int ile = lev[v] - lg + 1;
                wyn += ll(ile) * ll(zapytanie(0, m - 1, 0, kol[g].front().st - 1, 1));
                aktualizuj(0, m - 1, kol[g].front().st, ile, 1);
                operac.pb({kol[g].front().st, ile});
                kol[g].front().nd -= ile;
            }

            reverse(jakie.begin(), jakie.end());
            for (auto x : jakie) {
                wyn += ll(x.nd) * ll(zapytanie(0, m - 1, 0, x.st - 1, 1));
                aktualizuj(0, m - 1, x.st, x.nd, 1);
            }

            kol[g].push_front({c[b], lev[v] - lev[g] + 1});
            if (g == 1) {
                break;
            }
            v = ojc[g];
        }
        for (auto x : operac) {
            aktualizuj(0, m - 1, x.st, -x.nd, 1);
        }
        kol[gora[b]].pb({c[b], 1});

        cout << wyn << "\n";


    }

    return 0;
}


#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...