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>
//#pragma GCC optimize("O2")
using namespace std;
using ll = long long;
using ld = long double;
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define sz(x) (int)x.size()
#define endl '\n'
const int mod = 1e9 + 7;
const int inf = 2e9 + 5;
const ll linf = 9e18 + 5;
int n;
const int N = 1e5 + 5;
vector<int> adj[N];
int c[N];
int a[N];
int b[N];
int len[N];
int up[N];
int par[N];
int dtime = 0;
int st[N];
set<vector<int>> colors;
void dfssort(int u) {
len[u] = 1;
for (auto v : adj[u]) {
dfssort(v);
len[u] += len[v];
}
for (int i = 0; i < sz(adj[u]); i++) {
if (len[adj[u][0]] < len[adj[u][i]]) {
swap(adj[u][0], adj[u][i]);
}
}
}
void dfsuppar(int u) {
if (adj[u].empty()) {
return;
}
for (int i = 1; i < sz(adj[u]); i++) {
int v = adj[u][i];
up[v] = v;
par[v] = u;
dfsuppar(v);
}
up[adj[u][0]] = up[u];
par[adj[u][0]] = u;
dfsuppar(adj[u][0]);
}
void dfsst(int u) {
st[u] = dtime;
dtime++;
for (auto v : adj[u]) {
dfsst(v);
}
}
void init() {
}
void input() {
cin >> n;
for (int i = 1; i <= n; i++) {
cin >> c[i];
}
for (int i = 1; i < n; i++) {
cin >> a[i] >> b[i];
}
for (int i = 1; i < n; i++) {
adj[a[i]].pb(b[i]);
}
}
vector<int> cover(int x) {
return *prev(colors.upper_bound({x, inf, inf}));
}
void cut(int x) {
vector<int> v = cover(x);
colors.erase(v);
colors.insert({v[0], x, v[2]});
if (v[1] >= x + 1) {
colors.insert({x + 1, v[1], v[2]});
}
}
ll mrg(vector<pair<int, int>> &v, int l, int r, int m) {
ll ans = 0;
int sum = 0;
int i = l;
int j = m + 1;
int k = 0;
vector<pair<int, int>> sorted(r - l + 1);
while (i <= m && j <= r) {
if (v[i] > v[j]) {
sorted[k] = v[j];
sum += v[j].se;
j++;
k++;
}
else {
sorted[k] = v[i];
ans += 1ll * sum * v[i].se;
i++;
k++;
}
}
while (i <= m) {
sorted[k] = v[i];
ans += 1ll * sum * v[i].se;
i++;
k++;
}
while (j <= r) {
sorted[k] = v[j];
sum += v[j].se;
j++;
k++;
}
for (int i = 0; i < (r - l + 1); i++) {
v[l + i] = sorted[i];
}
return ans;
}
ll calcinvs(vector<pair<int, int>> &v, int l, int r) {
if (l == r) {
return 0;
}
int m = (l + r) / 2;
ll ans = 0;
ans += calcinvs(v, l, m);
ans += calcinvs(v, m + 1, r);
ans += mrg(v, l, r, m);
return ans;
}
ll calc(int vert) {
vector<pair<int, int>> cols;
int u = vert;
while (u) {
int v = up[u];
cut(st[v] - 1);
cut(st[u]);
int now = st[u];
while (st[v] <= now) {
vector<int> vect = cover(now);
cols.pb(mp(vect[2], vect[1] - vect[0] + 1));
colors.erase(vect);
now = vect[0] - 1;
}
colors.insert({st[v], st[u], c[vert]});
u = par[up[u]];
}
reverse(cols.begin(), cols.end());
// return calcinvs(cols, 0, sz(cols) - 1);
ll ans = 0;
for (int i = 0; i < sz(cols); i++) {
for (int j = i + 1; j < sz(cols); j++) {
if (cols[i].fi > cols[j].fi) {
ans += 1ll * cols[i].se * cols[j].se;
}
}
}
return ans;
}
void solve() {
dfssort(1);
up[1] = 1;
dfsuppar(1);
dfsst(1);
colors.insert({-1, n, inf});
calc(1);
for (int i = 1; i < n; i++) {
cout << calc(b[i]) << endl;
}
}
void output() {
}
int main() {
ios::sync_with_stdio(false);
cin.tie(0);
cout.tie(0);
int number_of_testcases = 1;
//cin >> number_of_testcases;
while (number_of_testcases--) {
init();
input();
solve();
output();
}
return 0;
}
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |