#include <bits/stdc++.h>
using namespace std;
template <typename T, int n, typename sum = std::plus<T> >
class Aib {
private:
T* aib;
T neutral;
sum adder;
inline int lsb(int x) {
return x & (-x);
}
public:
Aib(T _neutral) {
aib = new T[1+n];
for(int i = 1; i <= n; ++i)
aib[i] = neutral;
neutral = _neutral;
}
void update(int poz, T val) {
while(poz <= n) {
this->aib[poz] = adder(this->aib[poz], val);
poz += this->lsb(poz);
}
}
T query(int poz) {
T rez = this->aib[poz];
poz -= this->lsb(poz);
while(poz > 0) {
rez = adder(this->aib[poz], rez);
poz -= this->lsb(poz);
}
return rez;
}
};
const int MAX_N = 100000;
struct Cell {
int l, r, val;
};
vector<int> graph[1+MAX_N];
deque<Cell> chain[1+MAX_N];
deque<Cell> path;
int c[1+MAX_N];
int ma[MAX_N - 1], mb[MAX_N - 1];
int subarb[1+MAX_N], depth[1+MAX_N], parent[1+MAX_N];
int heavyId[1+MAX_N], head[1+MAX_N], lastid;
void dfs(int nod, int papa = 0) {
subarb[nod] = 1;
parent[nod] = papa;
for(auto it: graph[nod])
if(it != papa) {
depth[it] = depth[nod] + 1;
dfs(it, nod);
subarb[nod] += subarb[it];
}
}
void heavyTrav(int nod, int papa, int ch) {
int heavySon = -1;
heavyId[nod] = ++lastid;
head[nod] = ch;
chain[ch].push_back({depth[nod], depth[nod], c[nod]});
for(auto it: graph[nod])
if(it != papa && heavySon == -1)
heavySon = it;
else if(it != papa && subarb[it] > subarb[heavySon])
heavySon = it;
if(heavySon != -1)
heavyTrav(heavySon, nod, ch);
for(auto it: graph[nod])
if(it != papa && it != heavySon)
heavyTrav(it, nod, it);
}
void addRange(int nod, int l, int r, int val) {
while(chain[nod].size() > 0 && r >= chain[nod][0].r)
chain[nod].pop_front();
if(chain[nod].size() > 0 && r >= chain[nod][0].l)
chain[nod][0].l = r + 1;
chain[nod].push_front({l, r, val});
}
void update(int nod, int val) {
while(nod != 0) {
addRange(head[nod], depth[head[nod]], depth[nod], val);
nod = parent[head[nod]];
}
}
void extractRange(int nod, int r) {
int i = 0;
while(i + 1 < chain[nod].size() && chain[nod][i + 1].l <= r)
++i;
path.push_front({chain[nod][i].l, r, chain[nod][i].val});
--i;
while(i >= 0) {
path.push_front(chain[nod][i]);
--i;
}
}
void extract(int nod) {
while(nod != 0) {
extractRange(head[nod], depth[nod]);
nod = parent[head[nod]];
}
}
Aib<int, MAX_N> aib(0);
long long query(int nod) {
long long rez = 0LL;
path.clear();
extract(nod);
for(auto it: path) {
rez = rez + (long long)(it.r - it.l + 1) * aib.query(MAX_N) - aib.query(it.val);
aib.update(it.val, it.r - it.l + 1);
}
for(auto it: path)
aib.update(it.val, -(it.r - it.l + 1));
return rez;
}
int poz[1+MAX_N];
bool cmp(int a, int b) {
return c[a] < c[b];
}
void normalize(int n) {
for(int i = 1; i <= n; ++i)
poz[i - 1] = i;
std::sort(poz, poz + n, cmp);
int lastval = c[poz[0]], j = 1;
for(int i = 0; i < n; ++i)
if(c[poz[i]] == lastval)
c[poz[i]] = j;
else {
lastval = c[poz[i]];
c[poz[i]] = ++j;
}
}
int main() {
#ifdef HOME
FILE *fin = fopen("input.in", "r");
FILE *fout = fopen("output.out", "w");
#else
FILE *fin = stdin;
FILE *fout = stdout;
#endif
int n, a, b;
fscanf(fin, "%d", &n);
for(int i = 1; i <= n; ++i)
fscanf(fin, "%d", &c[i]);
normalize(n);
for(int i = 0; i < n - 1; ++i) {
fscanf(fin, "%d%d", &a, &b);
ma[i] = a;
mb[i] = b;
graph[a].push_back(b);
graph[b].push_back(a);
}
dfs(1);
heavyTrav(1, 0, 1);
for(int i = 0; i < n - 1; ++i) {
fprintf(fout, "%lld\n", query(ma[i]));
update(mb[i], c[mb[i]]);
}
#ifdef HOME
fclose(fin);
fclose(fout);
#endif
return 0;
}
Compilation message
construction.cpp: In function 'void extractRange(int, int)':
construction.cpp:98:14: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
while(i + 1 < chain[nod].size() && chain[nod][i + 1].l <= r)
~~~~~~^~~~~~~~~~~~~~~~~~~
construction.cpp: In function 'int main()':
construction.cpp:162:8: warning: ignoring return value of 'int fscanf(FILE*, const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
fscanf(fin, "%d", &n);
~~~~~~^~~~~~~~~~~~~~~
construction.cpp:164:9: warning: ignoring return value of 'int fscanf(FILE*, const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
fscanf(fin, "%d", &c[i]);
~~~~~~^~~~~~~~~~~~~~~~~~
construction.cpp:169:9: warning: ignoring return value of 'int fscanf(FILE*, const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
fscanf(fin, "%d%d", &a, &b);
~~~~~~^~~~~~~~~~~~~~~~~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
70 ms |
68856 KB |
Output is correct |
| 2 |
Incorrect |
81 ms |
68980 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
70 ms |
68856 KB |
Output is correct |
| 2 |
Incorrect |
81 ms |
68980 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
70 ms |
68856 KB |
Output is correct |
| 2 |
Incorrect |
81 ms |
68980 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
