/*
ID: varunra2
LANG: C++
TASK: highways
*/
#include <bits/stdc++.h>
using namespace std;
#ifdef DEBUG
#include "lib/debug.h"
#define debug(...) cerr << "[" << #__VA_ARGS__ << "]:", debug_out(__VA_ARGS__)
#define debug_arr(...) \
cerr << "[" << #__VA_ARGS__ << "]:", debug_arr(__VA_ARGS__)
#pragma GCC diagnostic ignored "-Wsign-compare"
//#pragma GCC diagnostic ignored "-Wunused-parameter"
//#pragma GCC diagnostic ignored "-Wunused-variable"
#else
#define debug(...) 42
#endif
#define EPS 1e-9
#define IN(A, B, C) assert(B <= A && A <= C)
#define INF (int)1e9
#define MEM(a, b) memset(a, (b), sizeof(a))
#define MOD 1000000007
#define MP make_pair
#define PB push_back
#define all(cont) cont.begin(), cont.end()
#define rall(cont) cont.end(), cont.begin()
#define x first
#define y second
const double PI = acos(-1.0);
typedef long long ll;
typedef long double ld;
typedef pair<int, int> PII;
typedef map<int, int> MPII;
typedef multiset<int> MSETI;
typedef set<int> SETI;
typedef set<string> SETS;
typedef vector<int> VI;
typedef vector<PII> VII;
typedef vector<VI> VVI;
typedef vector<string> VS;
#define rep(i, a, b) for (int i = a; i < (b); ++i)
#define trav(a, x) for (auto& a : x)
#define sz(x) (int)(x).size()
typedef pair<int, int> pii;
typedef vector<int> vi;
#pragma GCC diagnostic ignored "-Wsign-compare"
// util functions
int n;
VVI adj;
VI wch; // which chain
VI pos; // which position in the chain
VI sub; // subtree size
VI dep; // distance from the root
VI live; // liveliness
VII edgs;
void init() {
adj.resize(n);
wch.resize(n);
pos.resize(n);
sub.resize(n);
dep.resize(n);
live.resize(n);
edgs.resize(n - 1);
}
struct chain {
VII vals;
int ind;
int siz;
int par;
void init(int _ind, int _par) {
ind = _ind;
par = _par;
siz = 0;
}
void ins(int node) {
if (vals.empty())
vals.PB(MP(live[node], 1));
else if (vals.back().x == live[node])
vals.back().y++;
else
vals.PB(MP(live[node], 1));
siz++;
wch[node] = ind;
pos[node] = siz;
}
void done() {
// call this when you are done with this chain
reverse(all(vals));
}
VII get(int st, int val) {
// get from 0 -> st
VII ret;
int sum = 0;
while (!vals.empty()) {
auto x = vals.back();
int sum1 = sum + x.y;
if (st <= sum1) {
ret.PB(MP(x.x, st - sum));
vals.back().y -= (st - sum);
if (vals.back().y == 0) vals.pop_back();
break;
} else
ret.PB(x);
sum = sum1;
vals.pop_back();
}
vals.PB(MP(val, st));
return ret;
}
void deb() {
debug("new chain");
debug(ind);
debug(par);
debug(siz);
debug(vals);
debug("done");
}
};
vector<chain> chains;
struct BIT {
VI fen;
int n;
void init(int _n) {
n = _n;
fen.assign(n + 1, 0);
}
void upd(int ind, int val) {
ind++;
while (ind <= n) {
fen[ind] += val;
ind += (ind & -ind);
}
}
int qry(int ind) {
int ret = 0;
ind++;
while (ind >= 1) {
ret += fen[ind];
ind -= (ind & -ind);
}
return ret;
}
int qry(int l, int r) { return qry(r) - qry(l - 1); }
};
MPII cmp;
// first is val, second is count
void cmprs(VII vals) {
cmp.clear();
sort(all(vals));
vals.resize(unique(all(vals)) - vals.begin());
int cnt = 0;
trav(x, vals) cmp[x.x] = cnt++;
}
ll calcinvs(VII vals) {
ll ret = 0;
BIT bit;
bit.init(sz(vals));
cmprs(vals);
for (int i = 0; i < sz(vals); i++) {
ret += bit.qry(cmp[vals[i].x] + 1, sz(vals) - 1);
bit.upd(cmp[vals[i].x], vals[i].y);
}
return ret;
}
VII gen(int node, int og = -1) {
if (og == -1) og = node;
VII ret;
if (node == -1) return ret;
ret = gen(chains[wch[node]].par, og);
auto x = chains[wch[node]].get(pos[node], live[og]);
trav(xx, x) ret.PB(xx);
return ret;
}
ll qry(int node, int og) { return calcinvs(gen(node, og)); }
void dfs1(int u, int v) {
sub[u] = 1;
for (auto& x : adj[u]) {
if (x == v) continue;
dep[x] = dep[u] + 1;
dfs1(x, u);
sub[u] += sub[x];
}
}
void dfs2(int u, int v) {
if (sub[u] == 1) return;
if (u == 0) {
chain chn;
chn.init(0, -1);
chn.ins(0);
chains.PB(chn);
}
int mx = adj[u][0];
if (v == mx) mx = adj[u][1];
trav(x, adj[u]) {
if (x == v) continue;
if (sub[mx] < sub[x]) mx = x;
}
trav(x, adj[u]) {
if (x == v) continue;
if (x == mx) {
chains[wch[u]].ins(x);
} else {
chain nw;
nw.init(sz(chains), u);
nw.ins(x);
chains.PB(nw);
}
}
trav(x, adj[u]) {
if (x == v) continue;
dfs2(x, u);
}
}
int main() {
// #ifndef ONLINE_JUDGE
// freopen("highways.in", "r", stdin);
// freopen("highways.out", "w", stdout);
// #endif
cin.sync_with_stdio(0);
cin.tie(0);
cin >> n;
init();
for (int i = 0; i < n; i++) {
cin >> live[i];
}
for (int i = 0; i < n - 1; i++) {
int u, v;
cin >> u >> v;
u--, v--;
edgs[i] = MP(u, v);
adj[u].PB(v);
adj[v].PB(u);
}
dfs1(0, -1);
dfs2(0, -1);
trav(x, chains) x.done();
// debug(wch);
// debug(pos);
// trav(x, chains) x.deb();
for (int i = 0; i < n - 1; i++) {
cout << qry(edgs[i].x, edgs[i].y) << '\n';
}
// trav(x, chains) x.deb();
return 0;
}
Compilation message
construction.cpp: In member function 'void chain::deb()':
construction.cpp:19:20: warning: statement has no effect [-Wunused-value]
19 | #define debug(...) 42
| ^~
construction.cpp:124:5: note: in expansion of macro 'debug'
124 | debug("new chain");
| ^~~~~
construction.cpp:19:20: warning: statement has no effect [-Wunused-value]
19 | #define debug(...) 42
| ^~
construction.cpp:125:5: note: in expansion of macro 'debug'
125 | debug(ind);
| ^~~~~
construction.cpp:19:20: warning: statement has no effect [-Wunused-value]
19 | #define debug(...) 42
| ^~
construction.cpp:126:5: note: in expansion of macro 'debug'
126 | debug(par);
| ^~~~~
construction.cpp:19:20: warning: statement has no effect [-Wunused-value]
19 | #define debug(...) 42
| ^~
construction.cpp:127:5: note: in expansion of macro 'debug'
127 | debug(siz);
| ^~~~~
construction.cpp:19:20: warning: statement has no effect [-Wunused-value]
19 | #define debug(...) 42
| ^~
construction.cpp:128:5: note: in expansion of macro 'debug'
128 | debug(vals);
| ^~~~~
construction.cpp:19:20: warning: statement has no effect [-Wunused-value]
19 | #define debug(...) 42
| ^~
construction.cpp:129:5: note: in expansion of macro 'debug'
129 | debug("done");
| ^~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
1 ms |
364 KB |
Output is correct |
| 2 |
Correct |
1 ms |
364 KB |
Output is correct |
| 3 |
Incorrect |
1 ms |
364 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
1 ms |
364 KB |
Output is correct |
| 2 |
Correct |
1 ms |
364 KB |
Output is correct |
| 3 |
Incorrect |
1 ms |
364 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
1 ms |
364 KB |
Output is correct |
| 2 |
Correct |
1 ms |
364 KB |
Output is correct |
| 3 |
Incorrect |
1 ms |
364 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
