Submission #941967

#TimeUsernameProblemLanguageResultExecution timeMemory
941967juliany2Construction of Highway (JOI18_construction)C++17
100 / 100
895 ms23628 KiB
#include<bits/stdc++.h>
using namespace std;
using ll = long long;
#define all(x) (x).begin(), (x).end()

template<class T> struct BIT {
    int n;
    vector<T> bit;

    void init(int sz) { bit = vector<T> (n = sz + 7); }

    void upd(int i, T x) { 
        for (i++; i < n; i += i & -i) bit[i] += x;
    }

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

    T query(int l, int r) { return query(r) - query(l - 1); };
};

const int N = 1e5 + 7;
int n;
vector<int> adj[N];
array<int, 2> e[N];
int c[N], sz[N], head[N], pos[N], timer;
int parent[N];
set<array<int, 3>> s;
vector<array<int, 3>> cur;
BIT<int> bit;

void dfs(int v = 1, int p = 0) {
    parent[v] = p;
    sz[v] = 1;

    for (int &u : adj[v]) {
        if (u == p)
            continue;

        dfs(u, v);
        sz[v] += sz[u];

        if (adj[v][0] == p || sz[u] > sz[adj[v][0]])
            swap(u, adj[v][0]);
    }
}

void dfs_hld(int v = 1, int p = 0) {
    pos[v] = ++timer;

    for (int u : adj[v]) {
        if (u == p)
            continue;

        head[u] = (u == adj[v][0] ? head[v] : u);
        dfs_hld(u, v);
    }
}

void add(int l, int r, int x) {
    vector<array<int, 3>> tmp;
    while (s.lower_bound({l, -1, -1}) != s.end() && (*s.lower_bound({l, -1, -1}))[0] <= r) {
        auto [a, b, y] = *s.lower_bound({l, -1, -1});

        if (b <= r) {
            tmp.push_back({a, b, y});
            s.erase({a, b, y});
        }
        else {
            tmp.push_back({a, r, y});
            s.insert({r + 1, b, y});
            s.erase({a, b, y});
        }
    }

    s.insert({l, r, x});

    reverse(all(tmp));

    for (auto &[a, b, y] : tmp)
        cur.push_back({a, b, y});
}

void query(int v, int x) {
    for (; v >= 1; v = parent[head[v]])
        add(pos[head[v]], pos[v], x);
}

int main() {
    cin.tie(0)->sync_with_stdio(false);

    cin >> n;

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

    sort(all(cmp));
    cmp.erase(unique(all(cmp)), cmp.end());

    for (int i = 1; i <= n; i++)
        c[i] = lower_bound(all(cmp), c[i]) - cmp.begin() + 1;

    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;

        adj[u].push_back(v);
        adj[v].push_back(u);
        e[i] = {u, v};
    }

    dfs();

    head[1] = 1;
    dfs_hld();

    s.insert({1, 1, c[1]});

    bit.init(n + 1);

    for (int i = 1; i < n; i++) {
        auto [u, v] = e[i];

        query(v, c[v]);

        ll ans = 0;
        for (auto &[l, r, x] : cur) {
            int cnt = r - l + 1;
            ans += 1LL * bit.query(0, x - 1) * cnt;
            bit.upd(x, cnt);
        }

        for (auto &[l, r, x] : cur)
            bit.upd(x, -(r - l + 1));
        cur.clear();

        cout << ans << '\n';
    }

    return 0;
}

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