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>
#define FOR(i, x, n) for(int i = x; i < n; i++)
#define F0R(i, n) FOR(i, 0, n)
#define ROF(i, x, n) for(int i = n - 1; i >= x; i--)
#define R0F(i, n) ROF(i, 0, n)
#define WTF cout << "WTF" << endl
#define IOS ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0)
#define F first
#define S second
#define PB push_back
#define ALL(x) x.begin(), x.end()
#define RALL(x) x.rbegin(), x.rend()
using namespace std;
typedef long long LL;
typedef pair<int, int> PII;
typedef pair<LL, LL> PLL;
typedef vector<int> VI;
typedef vector<LL> VLL;
typedef vector<PII> VPII;
typedef vector<PLL> VPLL;
#define lc now << 1
#define rc now << 1 | 1
const int N = 1e5 + 5;
int n;
int cs[N], sz[N], t[N], cnt;
int head[N], height[N], par[N];
int order[N];
PII eset[N];
VI adj[N], comp;
struct Fenwick {
int fen[N];
VI changed;
void add(int id, int val) {
for(; id < N; id += id & -id) {
fen[id] += val;
changed.PB(id);
}
return;
}
int get(int id) {
int ret = 0;
for(; id; id -= id & -id) ret += fen[id];
return ret;
}
void clear() {
for(const int &on : changed) fen[on] = 0;
changed.clear();
}
} sum;
struct Segment {
bool tree[N << 2];
int lz[N << 2], cnt = 0;
void init() {
memset(lz, -1, sizeof lz);
return;
}
void shift(int now) {
if (lz[now] == -1) return;
tree[lc] = lz[lc] = lz[now];
tree[rc] = lz[rc] = lz[now];
lz[now] = -1;
return;
}
void change(int lq, int rq, bool val, int now = 1, int ls = 0, int rs = n - 1) {
if (rq < lq || rq < ls || rs < lq) return;
if (lq <= ls && rs <= rq) {
tree[now] = lz[now] = val;
return;
}
int mid = (ls + rs) >> 1;
shift(now);
change(lq, rq, val, lc, ls, mid);
change(lq, rq, val, rc, mid + 1, rs);
tree[now] = tree[lc] | tree[rc];
return;
}
int rightest1(int q, int now = 1, int ls = 0, int rs = n - 1) {
if (ls >= q || !tree[now]) return -1;
if (ls == rs) return ls;
int mid = (ls + rs) >> 1;
shift(now);
int case1 = rightest1(q, rc, mid + 1, rs);
return (case1 == -1 ? rightest1(q, lc, ls, mid) : case1);
}
} hldq;
struct MaxSeg {
PII tree[N << 2];
int cnt = 0;
void init() {
fill(tree, tree + (N << 2), PII(-1, -1));
return;
}
void cseg(int lq, int rq, PII val, int now = 1, int ls = 0, int rs = n - 1) {
if (rq < lq || rq < ls || rs < lq) return;
if (lq <= ls && rs <= rq) {
tree[now] = val;
return;
}
int mid = (ls + rs) >> 1;
cseg(lq, rq, val, lc, ls, mid);
cseg(lq, rq, val, rc, mid + 1, rs);
return;
}
void change(int lq, int rq, int val) {
cseg(lq, rq, PII(cnt, val));
cnt++;
}
PII gseg(int id, int now = 1, int ls = 0, int rs = n - 1) {
if (ls == rs) return tree[now];
int mid = (ls + rs) >> 1;
if (id <= mid) return max(tree[now], gseg(id, lc, ls, mid));
return max(tree[now], gseg(id, rc, mid + 1, rs));
}
int get(int id) {
return gseg(id).S;
}
} color;
void init() {
cin >> n;
comp.resize(n);
F0R(i, n) {
cin >> cs[i];
comp[i] = cs[i];
}
par[0] = -1;
sort(ALL(comp));
comp.resize(unique(ALL(comp)) - comp.begin());
F0R(i, n) cs[i] = lower_bound(ALL(comp), cs[i]) - comp.begin() + 1;
F0R(i, n - 1) {
int u, v; cin >> u >> v;
adj[--u].PB(--v);
eset[i] = {u, v};
par[v] = u;
}
return;
}
int preD(int now) {
sz[now] = 1;
for(int on : adj[now]) {
height[on] = height[now] + 1;
sz[now] += preD(on);
}
return sz[now];
}
void dfs(int now) {
order[cnt] = now;
t[now] = cnt++;
sort(ALL(adj[now]), [](const int &a, const int &b) {return sz[a] > sz[b];});
if (!adj[now].empty()) head[ adj[now][0] ] = head[now];
for(int on : adj[now])
dfs(on);
return;
}
LL get(int v) {
LL ret = 0;
sum.clear();
while(v != -1) {
int pl = hldq.rightest1(t[v]);
if (pl >= t[ head[v] ]) {
pl = order[pl];
LL us = height[v] - height[pl];
int num = color.get(t[v]);
ret += sum.get(num - 1) * us;
sum.add(num, us);
v = pl;
}
else {
LL us = height[v] - height[ head[v] ] + 1;
int num = color.get(t[v]);
ret += sum.get(num - 1) * us;
sum.add(num, us);
v = par[ head[v] ];
}
}
return ret;
}
void change(int v, int col) {
while(v != -1) {
hldq.change(t[v], t[v], 1);
//cout << "changing " << t[head[v]] << ' ' << t[v] << " to " << x
color.change(t[ head[v] ], t[v], col);
if (v == head[v]) {
v = par[v];
continue;
}
v = par[v];
hldq.change(t[ head[v] ], t[v], 0);
v = par[ head[v] ];
}
return;
}
int main() {
IOS;
init();
/*color.change(0, 0, 1);
color.change(0, 0, 2);
cout << color.get(0) << endl;*/
preD(0);
dfs(0);
iota(head, head + n, 0);
color.init();
color.change(0, 0, cs[0]);
hldq.init();
//F0R(i, n) cout << order[i] + 1 << ' '; cout << endl;
F0R(i, n - 1) {
cout << get(eset[i].F) << endl;
change(eset[i].S, cs[ eset[i].S ]);
//cout << "after edge " << i + 1 << " colors are ";
//F0R(i, n) cout << color.get(i) << ' '; cout << endl;
}
return 0;
}
Compilation message (stderr)
construction.cpp: In member function 'void Segment::shift(int)':
construction.cpp:78:27: warning: suggest parentheses around assignment used as truth value [-Wparentheses]
78 | tree[lc] = lz[lc] = lz[now];
| ~~~~~~~^~~~~~~~~
construction.cpp:79:27: warning: suggest parentheses around assignment used as truth value [-Wparentheses]
79 | tree[rc] = lz[rc] = lz[now];
| ~~~~~~~^~~~~~~~~
construction.cpp: In member function 'void Segment::change(int, int, bool, int, int, int)':
construction.cpp:88:33: warning: suggest parentheses around assignment used as truth value [-Wparentheses]
88 | tree[now] = lz[now] = val;
| ~~~~~~~~^~~~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |