#include <bits/stdc++.h>
using namespace std;
const int maxn = 1e5 + 5;
int lv[maxn], sz[maxn], par[maxn], tree[maxn], in[maxn], out[maxn], tin, sizeofchain[maxn];
map<int, int> mp;
vector<int> v[maxn], tocmp;
deque<pair<int, int>> vec[maxn];
int bigchild[maxn], chain[maxn], n, a[maxn], b[maxn], c[maxn];
void upd(int idx, int val) {
for (int i = idx; i < maxn; i+= i & (-i)) tree[i] += val;
}
int solve(int idx) {
int pas = 0;
for (int i = idx; i > 0; i -= i & (-i)) pas += tree[i];
return pas;
}
void dfs1(int a, int p) {
sz[a] = 1;
lv[a] = lv[p] + 1;
par[a] = p;
for (int i = 0; i < v[a].size(); i++) {
int to = v[a][i];
if (to == p) continue;
dfs1(to, a);
sz[a] += sz[to];
if (!bigchild[a] || sz[bigchild[a]] < sz[to]) bigchild[a] = to;
}
}
void dfs2(int a, int p) {
tin++;
in[a] = tin;
if (bigchild[a]) dfs2(bigchild[a], a);
for (int i = 0; i < v[a].size(); i++) {
int to = v[a][i];
if (to == p || to == bigchild[a]) continue;
dfs2(to, a);
}
out[a] = tin;
}
void chain_dfs(int a, int p) {
if (bigchild[a]) chain[bigchild[a]] = chain[a];
sizeofchain[a]++;
if (!vec[chain[a]].size()) vec[chain[a]].push_back({c[a], 1});
else {
if (vec[chain[a]].back().first == c[a]) vec[chain[a]].back().second++;
else vec[chain[a]].push_back({c[a], 1});
}
for (int i = 0; i < v[a].size(); i++) {
int to = v[a][i];
if (to == p) continue;
chain_dfs(to, a);
}
}
int solve(int a, int b) {
int res = 0;
int lst = chain[a];
int cur = a;
vector<pair<int, int>> pas;
while (true) {
lst = chain[cur];
int x = lv[cur] - lv[lst] + 1;
vector<pair<int, int>> vv;
int sum = 0;
for (int i = 0; i < vec[lst].size(); i++) {
if (vec[lst][i].second + sum > x) {
vv.push_back({vec[lst][i].first, (x - sum)});
break;
}
sum += vec[lst][i].second;
vv.push_back({vec[lst][i].first, vec[lst][i].second});
if (sum == x) break;
}
reverse(vv.begin(), vv.end());
for (int i = 0; i < vv.size(); i++) pas.push_back({vv[i].first, vv[i].second});
cur = par[lst];
if (cur == 0) break;
}
for (int i = 0; i < pas.size(); i++) {
res += solve(pas[i].first - 1) * pas[i].second;
upd(pas[i].first, pas[i].second);
}
for (int i = 0; i < pas.size(); i++) upd(pas[i].first, -pas[i].second);
return res;
}
void toadd(int a, int b) {
int lst = chain[a];
int cur = a;
while (true) {
lst = chain[cur];
int x = lv[cur] - lv[lst] + 1;
int sum = 0;
vector<pair<int, int>> vv;
while (vec[lst].size()) {
int i = 0;
if (vec[lst][i].second + sum > x) {
vec[lst][i].second = vec[lst][i].second - (x - sum);
break;
}
sum += vec[lst][i].second;
vec[lst].pop_front();
if (sum == x) break;
}
vec[lst].push_front({c[b], x});
cur = par[lst];
if (cur == 0) break;
}
}
int main() {
ios_base::sync_with_stdio(false);
cin.tie(nullptr);
cin >> n;
for (int i = 1; i <= n; i++) {
cin >> c[i];
tocmp.push_back(c[i]);
}
sort(tocmp.begin(), tocmp.end());
int cnt = 1;
mp[tocmp[0]] = 1;
for (int i = 1; i < tocmp.size(); i++) {
if (tocmp[i] != tocmp[i - 1]) cnt++;
mp[tocmp[i]] = cnt;
}
for (int i = 1; i <= n; i++) c[i] = mp[c[i]];
for (int i = 1; i <= n - 1; i++) {
cin >> a[i] >> b[i];
v[a[i]].push_back(b[i]);
v[b[i]].push_back(a[i]);
}
for (int i = 1; i <= n; i++) chain[i] = i;
dfs1(1, 0);
dfs2(1, 0);
chain_dfs(1, 0);
for (int i = 1; i <= n - 1; i++) {
int ans = solve(a[i], b[i]);
cout << ans << '\n';
}
}
Compilation message
construction.cpp: In function 'void dfs1(int, int)':
construction.cpp:26:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
26 | for (int i = 0; i < v[a].size(); i++) {
| ~~^~~~~~~~~~~~~
construction.cpp: In function 'void dfs2(int, int)':
construction.cpp:39:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
39 | for (int i = 0; i < v[a].size(); i++) {
| ~~^~~~~~~~~~~~~
construction.cpp: In function 'void chain_dfs(int, int)':
construction.cpp:55:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
55 | for (int i = 0; i < v[a].size(); i++) {
| ~~^~~~~~~~~~~~~
construction.cpp: In function 'int solve(int, int)':
construction.cpp:72:21: warning: comparison of integer expressions of different signedness: 'int' and 'std::deque<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
72 | for (int i = 0; i < vec[lst].size(); i++) {
| ~~^~~~~~~~~~~~~~~~~
construction.cpp:82:21: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
82 | for (int i = 0; i < vv.size(); i++) pas.push_back({vv[i].first, vv[i].second});
| ~~^~~~~~~~~~~
construction.cpp:86:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
86 | for (int i = 0; i < pas.size(); i++) {
| ~~^~~~~~~~~~~~
construction.cpp:90:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
90 | for (int i = 0; i < pas.size(); i++) upd(pas[i].first, -pas[i].second);
| ~~^~~~~~~~~~~~
construction.cpp: In function 'int main()':
construction.cpp:129:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
129 | for (int i = 1; i < tocmp.size(); i++) {
| ~~^~~~~~~~~~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
40 ms |
69972 KB |
Output is correct |
| 2 |
Incorrect |
43 ms |
70024 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
40 ms |
69972 KB |
Output is correct |
| 2 |
Incorrect |
43 ms |
70024 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
40 ms |
69972 KB |
Output is correct |
| 2 |
Incorrect |
43 ms |
70024 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
