Submission #1256549

#TimeUsernameProblemLanguageResultExecution timeMemory
1256549nguynConstruction of Highway (JOI18_construction)C++20
100 / 100
938 ms18632 KiB
#include <bits/stdc++.h>
using namespace std;

#define ll long long
#define F first
#define S second
#define pb push_back
#define pii pair<int, int>

const int N = 1e5 + 5;
const int mod = 1e9 + 7;

int n, a[N];
int sz[N], par[N], head[N], pos[N], timedfs = 0;
pii edges[N];
vector<int> g[N];
vector<pii> vec[N];
vector<int> compress;
int vis[N];
int bit[N];

void reset(int id) {
    for (int i = id; i <= n; i += (i & -i)) bit[i] = 0;
}

void update(int id, int val) {
    for (int i = id; i <= n; i += (i & -i)) bit[i] += val;
}

int get(int id) {
    int ret = 0;
    for (int i = id; i > 0; i -= (i & -i)) ret += bit[i];
    return ret;
}

void dfs(int u, int p) {
    sz[u] = 1;
    for (int v : g[u]) {
        if (v == p) continue;
        par[v] = u;
        dfs(v, u);
        sz[u] += sz[v];
    }
}

void hld(int u, int cur_head, int p) {
    pos[u] = ++timedfs;
    head[u] = cur_head;
    int big = 0;
    for (int v : g[u]) {
        if (v == p) continue;
        if (sz[big] < sz[v]) big = v;
    }
    if (big) hld(big, cur_head, u);
    for (int v : g[u]) {
        if (v == p || v == big) continue;
        hld(v, v, u);
    }
}

void upd(int u, int x) {
    while(u != 0) {
        while(!vec[head[u]].empty() && pos[vec[head[u]].back().F] <= pos[u]) vec[head[u]].pop_back();
        vec[head[u]].pb({u, x});
        u = par[head[u]];
    }
}

ll query(int u) {
    ll ret = 0;
    vector<int> tmp;
    while(u != 0) {
        vector<pii>& cur = vec[head[u]];
        if (!cur.empty()) {
            int j = (int)cur.size() - 1;
            for (int i = 0; i < (int)cur.size() - 1; i++) {
                if (pos[cur[i + 1].F] <= pos[u]) {
                    j = i;
                    break;
                }
            }
            for (int i = j; i < cur.size(); i++) {
                int cur_pos = min(pos[cur[i].F], pos[u]);
                int cnt = (i != (int)cur.size() - 1 ? cur_pos - pos[cur[i + 1].F] : cur_pos - pos[head[u]] + 1);
                ret += 1ll * cnt * get(cur[i].S - 1);
                update(cur[i].S, cnt);
                if (!vis[cur[i].S]) {
                    vis[cur[i].S] = 1;
                    tmp.pb(cur[i].S);
                }
            }
        }
        u = par[head[u]];
    }
    for (auto x : tmp) {
        vis[x] = 0;
        reset(x);
    }
    return ret;
}

signed main() {
    ios_base::sync_with_stdio(0);
    cin >> n;
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
        compress.pb(a[i]);
    }
    sort(compress.begin(), compress.end());
    compress.erase(unique(compress.begin(), compress.end()), compress.end());
    for (int i = 1; i <= n; i++) {
        a[i] = lower_bound(compress.begin(), compress.end(), a[i]) - compress.begin() + 1;
    }
    for (int i = 1; i < n; i++) {
        int u, v; cin >> u >> v;
        edges[i] = {u, v};
        g[u].pb(v);
        g[v].pb(u);
    }
    dfs(1, 0);
    hld(1, 1, 0);
    for (int i = 1; i < n; i++) {
        cout << query(edges[i].F) << '\n';
        upd(edges[i].S, a[edges[i].S]);
    }
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...