Submission #1031566

#TimeUsernameProblemLanguageResultExecution timeMemory
1031566underwaterkillerwhaleConstruction of Highway (JOI18_construction)C++17
100 / 100
138 ms20312 KiB
#include <bits/stdc++.h>
#define se              second
#define fs              first
#define mp              make_pair
#define pb              push_back
#define ll              long long
#define ii              pair<ll,ll>
#define ld              long double
#define SZ(v)           (int)v.size()
#define ALL(v)          v.begin(), v.end()
#define bit(msk, i)     ((msk >> i) & 1)
#define iter(id, v)     for(auto id : v)
#define rep(i,m,n)      for(int i=(m); i<=(n); i++)
#define reb(i,m,n)      for(int i=(m); i>=(n); i--)
using namespace std;

mt19937_64 rd(chrono :: steady_clock :: now().time_since_epoch().count());
ll Rand(ll l, ll r) { return uniform_int_distribution<ll> (l, r)(rd); }

const int N = 1e5 + 7;
const int Mod = 1e9 + 7;
const int szBL = 320;
const ll INF = 1e9 + 160907;
const int BASE = 137;

int n, m;
int a[N];
vector<int> ke[N];
vector<pair<int,int>> vec[N];
pair<int,int> qr[N];

struct fenwick_Tree {
    int m;
    int fen[N];
    void init (int n) {
        m = n;
        rep (i, 1, m) fen[i] = 0;
    }
    void update (int pos, int val) {
        for (; pos <= m; pos += pos & -pos) fen[pos] += val;
    }
    int get (int pos) {
        int res = 0;
        for (;pos > 0; pos -= pos & -pos) res += fen[pos];
        return res;
    }
}fen;

int cid[N], Head[N], par[N], high[N], nChild[N], numC = 1, pin[N];

void pdfs (int u, int p) {
    par[u] = p;
    nChild[u] = 1;
    high[u] = high[p] + 1;
    iter (&v, ke[u]) if (v != p){
        pdfs(v, u);
        nChild[u] += nChild[v];
    }
}

void hld (int u, int p) {
    static int time_dfs = 0;
    Head[u] = p;
    pin[u] = ++time_dfs;
    cid[u] = numC;
    int mxV = -1;
    iter (&v, ke[u])
        if (pin[v] == 0 && (mxV == -1 || nChild[v] > nChild[mxV])) mxV = v;
    if (mxV != -1) hld (mxV, p);
    iter (&v, ke[u]) {
        if (pin[v] == 0) {
            ++numC;
            hld(v, v);
        }
    }
}

ll update (int u, int col) {
    static vector<pair<int,int>> chains; chains.clear();
    static vector<pair<int,int>> del; del.clear();
    static vector<int> vals; vals.clear();
    for (; u != 0; u = par[Head[u]]) chains.push_back({cid[u], high[u] - high[Head[u]] + 1});
    reverse(ALL(chains));
    iter (&ch, chains) {
        vector<pair<int,int>> &cur = vec[ch.fs];
        int X = ch.se;
        int pre = 0;
        while (!cur.empty()) {
            pair<int, int> &u = cur.back();
            if (pre + u.se > X) {
                u.se = u.se - (X - pre);
                vals.push_back(u.fs);
                del.push_back({u.fs, X - pre});
                break;
            }
            else {
                del.push_back(u);
                vals.push_back(u.fs);
                pre += u.se;
                cur.pop_back();
            }
        }
        cur.push_back({col, X});
    }
    sort (ALL(vals));
    int ntot; vals.resize(ntot = unique(ALL(vals)) - vals.begin());
    fen.init(ntot);
    ll res = 0;
    iter (&cur, del) {
        int pos = lower_bound(ALL(vals), cur.fs) - vals.begin() + 1;
        res += 1LL * (fen.get(ntot) - fen.get(pos)) * cur.se;
        fen.update(pos, cur.se);
    }
    return res;
}

void solution () {
    cin >> n;
    rep (i, 1, n) cin >> a[i];
    rep (i, 1, n - 1) {
        int u, v;
        cin >> u >> v;
        qr[i] = {u, v};
        ke[u].push_back(v);
        ke[v].push_back(u);
    }
    pdfs(1, 0);
    hld(1, 1);
    vec[1].push_back({a[1], 1});
    rep (i, 1, n - 1) {
        int u, v;
        tie(u, v) = qr[i];
        cout << update (v, a[v]) <<"\n";
    }
}

#define file(name) freopen(name".inp", "r", stdin); \
freopen(name".out", "w", stdout);

int main () {
//    file("c");
    ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    ll num_Test = 1;
//    cin >> num_Test;
    while(num_Test--)
        solution();
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...