This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include <bits/stdc++.h>
using namespace std;
#define cerr if (0) cerr
#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;
}
cerr << "size(" << i+1 << ") => " << s << endl;
return s;
}
void dfs_hvy(int r, int i) {
cerr << "dfs_hvy(" << r+1 << ", " << i+1 << ")\n";
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) {
cerr << "heavy chain: ";
for (int i = 0; i < sz(v); i++) {
s.insert({i, i, c[v[i]]});
cerr << v[i] + 1 << " \n"[i == sz(v) - 1];
}
}
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);
s.insert(two);
return s.find(two);
}
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 i = 0; i < n; i++) cerr << c[i] << " \n"[i == n - 1];
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) {
cerr << "heavy root: " << i+1 << '\n';
ds[i] = DS(chain[i]);
}
for (auto [i, j] : con) {
vector<pii> v;
for (auto [r, k] : path(i)) ds[r].qry(k, v);
for (auto [c, k] : v) cerr << "(color=" << c << ", count=" << k << "), ";
cerr << endl;
cout << num_inversions(v) << '\n';
for (auto [r, k] : path(j)) ds[r].color(k, c[j]);
}
}
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |