Submission #1283806

#TimeUsernameProblemLanguageResultExecution timeMemory
1283806IcelastConstruction of Highway (JOI18_construction)C++20
100 / 100
235 ms65584 KiB
#include <iostream>
#include <bits/stdc++.h>
#define ll long long
using namespace std;
const ll maxn = 2*1e5+5, INF = 4e18+9;
template <class T>
struct Fenwick {
    int n, log;
    vector<T> bit;
    Fenwick(){}
    Fenwick(int n) : n(n), log(32 - __builtin_clz(n + 1)), bit(n + 1, 0) {}
    void add(int i, T delta) {
        for (; i <= n; i += i & -i) {
            bit[i] += delta;
        }
    }

    T sum(int i) {
        T res = 0;
        for (; i > 0; i -= i & -i) {
            res += bit[i];
        }
        return res;
    }
    T sum(int l, int r) {
        return sum(r)-sum(l-1);
    }
    int kth(T k) {
        T sum = 0;
        int pos = 0;
        for (int l = log - 1; l >= 0; l--) {
            if (pos + (1 << l) <= n && sum + bit[pos + (1 << l)] <= k) {
                pos += 1 << l;
                sum += bit[pos];
            }
        }
        return pos;
    }
};
struct HeavyLightDecomposition{
    int n, N, root;
    vector<int> sz, pa, head, depth, lt, rt;
    vector<vector<pair<int, int>>> split;
    Fenwick<int> bit;
    HeavyLightDecomposition(int n, int N, int root, vector<vector<int>> &adj) : n(n), N(N), root(root){
        sz.resize(n+1);
        pa.resize(n+1, -1);
        head.resize(n+1);
        depth.resize(n+1, -1);
        lt.resize(n+1);
        rt.resize(n+1);
        split.resize(n+1);
        bit = Fenwick<int>(N+1);
        auto root_tree = [&](auto self, int u, int p) -> void{
            sz[u] = 1;
            depth[u] = depth[p] + 1;
            pa[u] = p;
            for(int v : adj[u]){
                if(v == p) continue;
                self(self, v, u);
                sz[u] += sz[v];
            }
        };
        root_tree(root_tree, 1, 0);
        int timer = 0;
        auto decompose = [&](auto decompose, int u, int h) -> void{
            timer++;
            lt[u] = timer;
            head[u] = h;
            int heavy = -1;
            for(int v : adj[u]){
                if(v == pa[u]) continue;
                if(heavy == -1 || sz[heavy] < sz[v]) heavy = v;
            }
            if(heavy != -1){
                decompose(decompose, heavy, h);
            }
            for(int v : adj[u]){
                if(v == pa[u] || v == heavy) continue;
                decompose(decompose, v, v);
            }
            rt[u] = timer;
        };
        decompose(decompose, root, root);
    };
    int lca(int u, int v){
        for(; head[u] != head[v]; v = pa[head[v]]){
            if(depth[head[u]] > depth[head[v]]) swap(u, v);
        }
        if(depth[u] > depth[v]) swap(u, v);
        return u;
    }
    int dist(int u, int v) {
        return depth[u]+depth[v]-2*depth[lca(u, v)];
    }
    ll work(int v, int x){
        ll ans = 0;
        vector<pair<int, int>> to_add;
        vector<pair<int, int>> roll_back;
        while(v != 0){
            to_add.clear();
            int h = head[v];
            to_add.push_back({lt[h]-1, -1});
            while(!split[h].empty() && split[h].back().first <= lt[v]){
                to_add.push_back(split[h].back());
                split[h].pop_back();

            }
            if(!split[h].empty()) to_add.push_back({lt[v], split[h].back().second});
            for(int i = (int)to_add.size()-1; i >= 1; i--){
                int len = to_add[i].first - to_add[i-1].first;
                int c = to_add[i].second;
                ans += bit.sum(1, c-1) * len;
                roll_back.push_back({c, len});
                bit.add(c, len);
            }
            split[h].push_back({lt[v], x});
            v = pa[head[v]];
        }
        for(auto it : roll_back){
            bit.add(it.first, -it.second);
        }
        return ans;
    }
};
struct edge{
    int u, v;
};
struct normalize{
    vector<ll> poi, pot;
    void add(ll x){
        poi.push_back(x);
    }
    void start(){
        sort(poi.begin(), poi.end());
        if(poi.size() > 0) pot.push_back(poi[0]);
        for(int i = 1; i < (int)poi.size(); i++){
            if(poi[i] != poi[i-1]){
                pot.push_back(poi[i]);
            }
        }
    }
    int encode(ll x){
        return lower_bound(pot.begin(), pot.end(), x) - pot.begin()+1;
    }
};
void solve(){
    int n;
    cin >> n;
    vector<int> a(n+1);
    for(int i = 1; i <= n; i++){
        cin >> a[i];
    }
    vector<vector<int>> adj(n+1);
    vector<edge> e;
    for(int i = 1; i < n; i++){
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
        e.push_back({u, v});
    }
    normalize norm;
    for(int i = 1; i <= n; i++){
        norm.add(a[i]);
    }
    norm.start();
    for(int i = 1; i <= n; i++){
        a[i] = norm.encode(a[i]);
    }
    int N = norm.pot.size();
    HeavyLightDecomposition HLD(n, N, 1, adj);
    HLD.work(1, a[1]);
    for(int i = 1; i < n; i++){
        int u = e[i-1].u, v = e[i-1].v;
        ll res = HLD.work(v, a[v]);
        cout << res << "\n";
    }
}
int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    solve();
}

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