#include <bits/stdc++.h> // PLEASE
using namespace std;
typedef long long ll;
typedef long double ld;
typedef pair <int, int> pp;
#define MAXN 100005
#define MAXM 1005
#define MAXP 25
#define A first
#define B second
#define MP make_pair
#define PB push_back
#define PF push_front
const ll INF = 2e9+13;
const ll MOD = 1e9+7;
int N;
int C[MAXN];
vector <int> g[MAXN];
int in[MAXN], rin[MAXN], top[MAXN], out[MAXN], sz[MAXN], d[MAXN], par[MAXN];
int c = 1;
void dfs1(int u)
{
sz[u] = 1;
for(auto v : g[u]) {
d[v] = d[u] + 1; dfs1(v); sz[u] += sz[v];
if(sz[v] > sz[g[u][0]]) swap(v, g[u][0]);
}
}
void dfs2(int u)
{
in[u] = c++;
rin[in[u]] = u;
for(auto v : g[u]) {
if(v == g[u][0]) top[v] = top[u];
else top[v] = v;
dfs2(v);
} out[u] = c;
}
struct D {
int num, cnt, dep;
};
struct BIT { // 1-indexed
ll T[MAXN];
void init() { memset(T, 0, sizeof(T)); }
void upd(int x, ll v) { while(x < MAXN) { T[x] += v; x += (x & -x); } }
ll qry(int x) { ll ret = 0; while(x > 0) { ret += T[x]; x -= (x & -x); } return ret; }
ll get(int a, int b) {
if(a > b) return 0;
return qry(b) - qry(a-1);
}
} FEN;
deque <D> Q[MAXN];
vector <pp> ups;
vector <D> enc;
ll qry(int u)
{
vector <int> todo;
int tmp = u;
while(top[tmp] != 1) {
tmp = par[top[tmp]];
todo.PB(top[tmp]);
}
reverse(todo.begin(), todo.end());
ll ret = 0;
ups.clear(); enc.clear();
for(int i=0; i<todo.size(); i++) {
int root = todo[i];
int nxt = top[u]; if(i+1 < todo.size()) nxt = todo[i+1];
int cn = 0;
for(auto V : Q[root]) {
cn += V.cnt; ll cv = V.cnt;
if(cn >= d[nxt] - d[root]) cv -= (ll)(cn - (d[nxt] - d[root]));
ret += FEN.get(V.num + 1, MAXN-1) * cv;
FEN.upd(V.num, cv); ups.PB({V.num, cv});
if(cn >= d[nxt] - d[root]) break;
}
}
// now for the remaining
int cn = 0;
for(auto V : Q[top[u]]) {
cn += V.cnt;
ll cv = V.cnt;
if(cn >= d[u] - d[top[u]] + 1) cv -= (ll)(cn - (d[u] - d[top[u]] + 1));
ret += FEN.get(V.num + 1, MAXN-1) * cv;
FEN.upd(V.num, cv); ups.PB({V.num, cv});
if(cn >= d[u] - d[top[u]] + 1) break;
}
for(auto up : ups) FEN.upd(up.A, -up.B);
return ret;
}
void kill(int u, int col, int add)
{
int t = top[u];
int cn = 0; bool pushed = 0;
while(Q[t].size() && (Q[t].front().dep <= d[u])) {
cn += Q[t].front().cnt;
D tmp = Q[t].front();
Q[t].pop_front();
if(cn >= d[u] - d[t] + 1) { // split this one
cn -= tmp.cnt;
cn += d[u] - tmp.dep + 1;
tmp.cnt -= d[u] - tmp.dep + 1;
if(tmp.cnt) Q[t].PF({tmp.num, tmp.cnt, d[t] + cn});
cn += add; pushed = 1;
Q[t].PF({col, cn, d[t]}); cn = 0;
break;
}
}
if(!pushed) Q[t].PF({col, cn+add, d[t]});
}
void upd(int u)
{
int col = C[u];
// update this chain
kill(u, col, 1);
// now kill all other chains
while(top[u] != 1) {
u = par[top[u]];
kill(u, col, 0);
}
/*cout << endl;
cout << "PRINTING" << endl;
for(int i=1; i<=N; i++) {
if(i == top[i]) {
cout << "from node " << i << ":" << endl;
for(auto V : Q[i]) cout << V.num << ' ' << V.cnt << ' ' << V.dep << endl;
}
}*/
}
int main(void)
{
ios_base::sync_with_stdio(false);
cin.tie(0);
cin >> N; vector <int> v;
for(int i=1; i<=N; i++) { cin >> C[i]; v.PB(C[i]); }
sort(v.begin(), v.end()); v.resize(unique(v.begin(), v.end()) - v.begin());
for(int i=1; i<=N; i++) C[i] = lower_bound(v.begin(), v.end(), C[i]) - v.begin() + 1;
vector <pp> ord;
for(int i=0; i<N-1; i++) {
int p, u; cin >> p >> u;
par[u] = p;
g[p].PB(u);
ord.PB({p, u});
}
par[1] = 1; top[1] = 1; dfs1(1); dfs2(1); Q[1].PF({C[1], 1, d[1]});
FEN.init();
for(auto q : ord) {
int p = q.A; int u = q.B;
//cout << "ANSWER ";
cout << qry(p) << endl;
upd(u);
}
}
Compilation message
construction.cpp: In function 'll qry(int)':
construction.cpp:70:16: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
for(int i=0; i<todo.size(); i++) {
~^~~~~~~~~~~~
construction.cpp:72:28: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
int nxt = top[u]; if(i+1 < todo.size()) nxt = todo[i+1];
~~~~^~~~~~~~~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
71 ms |
69240 KB |
Output is correct |
| 2 |
Execution timed out |
1084 ms |
69348 KB |
Time limit exceeded |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
71 ms |
69240 KB |
Output is correct |
| 2 |
Execution timed out |
1084 ms |
69348 KB |
Time limit exceeded |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
71 ms |
69240 KB |
Output is correct |
| 2 |
Execution timed out |
1084 ms |
69348 KB |
Time limit exceeded |
| 3 |
Halted |
0 ms |
0 KB |
- |
