제출 #892773

#제출 시각아이디문제언어결과실행 시간메모리
892773boxConstruction of Highway (JOI18_construction)C++17
100 / 100
413 ms37408 KiB
#include <bits/stdc++.h>
using namespace std;

#define ar array
#define sz(v) int(std::size(v))
using i64 = long long;
using pii = pair<int, int>;

const int N = 1e5;

int n, c[N];
vector<int> adj[N];
vector<pii> con;
vector<int> chain[N];
int parent[N], child[N], root[N], depth[N];

int dfs_sz(int i) {
    int s = 1, x = 0;
    child[i] = -1;
    for (int j : adj[i]) if (parent[i] != j) {
        parent[j] = i;
        depth[j] = depth[i] + 1;
        int y = dfs_sz(j);
        if (y > x) {
            x = y;
            child[i] = j;
        }
        s += y;
    }
    return s;
}

void dfs_hvy(int r, int i) {
    root[i] = r;
    chain[r].push_back(i);
    if (~child[i]) dfs_hvy(r, child[i]);
    for (int j : adj[i]) if (parent[i] != j && child[i] != j) dfs_hvy(j, j);
}

struct node {
    int l, r, c;
    bool operator<(const node o) const {
        return r < o.r;
    }
};

struct DS {
    set<node> s;
    DS() {}
    DS(vector<int> v) {
        for (int i = 0; i < sz(v); i++)
            s.insert({i, i, c[v[i]]});
    }
    auto cut(int k) {
        auto it = s.lower_bound({0, k, 0});
        if (it == end(s) || it->l == k) return it;
        node one{it->l, k - 1, it->c};
        node two{k, it->r, it->c};
        s.erase(it);
        s.insert(one);
        return s.insert(two).first;
    }
    void color(int k, int c) {
        auto r = cut(k);
        s.erase(begin(s), r);
        s.insert({0, k - 1, c});
    }
    void qry(int k, vector<pii> &v) {
        auto r = cut(k);
        for (auto l = begin(s); l != r; ++l) v.push_back({l->c, l->r - l->l + 1});
    }
} ds[N];

struct {
    int t[N];
    void add(int i, int x) {
        while (i < N) {
            t[i] += x;
            i |= i + 1;
        }
    }
    int qry(int i) {
        int x = 0;
        while (i >= 0) {
            x += t[i];
            i &= i + 1, i--;
        }
        return x;
    }
} ft;

i64 num_inversions(vector<pii> v) {
    i64 num = 0;
    reverse(begin(v), end(v));
    for (auto [c, k] : v) {
        num += (i64) ft.qry(c - 1) * k;
        ft.add(c, k);
    }
    for (auto [c, k] : v) ft.add(c, -k);
    return num;
}

vector<pii> path(int i) {
    vector<pii> v;
    while (~i) {
        v.push_back({root[i], depth[i] - depth[root[i]] + 1});
        i = parent[root[i]];
    }
    reverse(begin(v), end(v));
    return v;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cin >> n;
    for (int i = 0; i < n; i++) cin >> c[i];
    auto sc = set(c, c + n);
    auto vc = vector(begin(sc), end(sc));
    for (int i = 0; i < n; i++) c[i] = lower_bound(begin(vc), end(vc), c[i]) - begin(vc);
    for (int h = 0; h < n - 1; h++) {
        int i, j; 
        cin >> i >> j, i--, j--;
        adj[i].push_back(j);
        adj[j].push_back(i);
        con.push_back({i, j});
    }
    parent[0] = -1;
    dfs_sz(0);
    dfs_hvy(0, 0);
    for (int i = 0; i < n; i++) if (root[i] == i) ds[i] = DS(chain[i]);
    for (auto [i, j] : con) {
        vector<pii> v;
        for (auto [r, k] : path(i)) ds[r].qry(k, v);
        cout << num_inversions(v) << '\n';
        for (auto [r, k] : path(j)) ds[r].color(k, c[j]);
    }
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...