#include <iostream>
#include <vector>
#include <chrono>
#include <random>
#include <limits>
#include <set>
#include <algorithm>
std::mt19937 rng((int) std::chrono::steady_clock::now().time_since_epoch().count());
class HLD {
public:
void init(int n) {
// this doesn't delete edges!
sz.resize(n);
in.resize(n);
out.resize(n);
rin.resize(n);
p.resize(n);
edges.resize(n);
nxt.resize(n);
h.resize(n);
}
void addEdge(int u, int v) {
edges[u].push_back(v);
edges[v].push_back(u);
}
void setRoot(int n) {
t = 0;
p[n] = n;
h[n] = 0;
prep(n, n);
nxt[n] = n;
hld(n);
}
int getLCA(int u, int v) {
while(!inSubtree(nxt[u], v)) {
u = p[nxt[u]];
}
while(!inSubtree(nxt[v], u)) {
v = p[nxt[v]];
}
return in[u] < in[v] ? u : v;
}
bool inSubtree(int u, int v) {
// is v in the subtree of u?
return in[u] <= in[v] && in[v] < out[u];
}
std::vector<std::pair<int, int>> getPathtoAncestor(int u, int anc) {
// returns ranges [l, r) that the path has
std::vector<std::pair<int, int>> ans;
//assert(inSubtree(anc, u));
while(nxt[u] != nxt[anc]) {
ans.emplace_back(in[nxt[u]], in[u] + 1);
u = p[nxt[u]];
}
// this includes the ancestor!
ans.emplace_back(in[anc], in[u] + 1);
return ans;
}
private:
std::vector<int> in, out, p, rin, sz, nxt, h;
std::vector<std::vector<int>> edges;
int t;
void prep(int on, int par) {
sz[on] = 1;
p[on] = par;
for(int i = 0; i < (int) edges[on].size(); i++) {
int &u = edges[on][i];
if(u == par) {
std::swap(u, edges[on].back());
edges[on].pop_back();
i--;
} else {
h[u] = 1 + h[on];
prep(u, on);
sz[on] += sz[u];
if(sz[u] > sz[edges[on][0]]) {
std::swap(edges[on][0], u);
}
}
}
}
void hld(int on) {
in[on] = t++;
rin[in[on]] = on;
for(auto u : edges[on]) {
nxt[u] = (u == edges[on][0] ? nxt[on] : u);
hld(u);
}
out[on] = t;
}
};
template <class Info = int>
class ColorUpdate {
public:
struct Range {
Range(int ll = 0) { this->l = ll; }
Range(int ll, int rr, Info vv) {
l = ll;
r = rr;
v = vv;
}
int l, r;
Info v;
bool operator < (const Range &b) const { return l < b.l; }
};
std::vector<Range> upd(int l, int r, Info v) {
std::vector<Range> ans;
if(l >= r) return ans;
auto it = ranges.lower_bound(l);
if(it != ranges.begin()) {
it--;
if(it->r > l) {
auto cur = *it;
ranges.erase(it);
ranges.insert(Range(cur.l, l, cur.v));
ranges.insert(Range(l, cur.r, cur.v));
}
}
it = ranges.lower_bound(r);
if(it != ranges.begin()) {
it--;
if(it->r > r) {
auto cur = *it;
ranges.erase(it);
ranges.insert(Range(cur.l, r, cur.v));
ranges.insert(Range(r, cur.r, cur.v));
}
}
for(it = ranges.lower_bound(l); it != ranges.end() && it->l < r; it++) {
ans.push_back(*it);
}
ranges.erase(ranges.lower_bound(l), ranges.lower_bound(r));
ranges.insert(Range(l, r, v));
return ans;
}
private:
std::set<Range> ranges;
};
template <class T>
class FenwickTree {
public:
void init(int s) {
n = s;
bit.assign(n + 1, 0);
}
void init(const std::vector<T> &a) {
n = a.size();
bit.assign(n + 1, 0);
for(int i = 1; i <= n; i++) {
bit[i] += a[i - 1];
if(i + (i & -i) <= n) {
bit[i + (i & -i)] += bit[i];
}
}
}
T qry(int x) {
x = std::min(x, (int)bit.size() - 1);
T ans = 0;
for(; x > 0; x -= x & -x) {
ans += bit[x];
}
return ans;
}
void upd(int x, T v) {
x = std::max(x, 1);
for(; x <= n; x += x & -x) {
bit[x] += v;
}
}
private:
int n;
std::vector<T> bit;
};
int main() {
std::ios_base::sync_with_stdio(false); std::cin.tie(NULL);
int n;
std::cin >> n;
std::vector<int> c(n);
for(int i = 0; i < n; i++) {
std::cin >> c[i];
}
auto vals = c;
std::sort(vals.begin(), vals.end());
for(auto &v : c) {
v = std::upper_bound(vals.begin(), vals.end(), v) - vals.begin();
}
std::vector<std::pair<int, int>> edges;
HLD tree;
tree.init(n);
for(int i = 1; i < n; i++) {
int u, v;
std::cin >> u >> v;
u--;v--;
edges.emplace_back(u, v);
tree.addEdge(u, v);
}
tree.setRoot(0);
ColorUpdate<int> colors;
FenwickTree<int> fw;
fw.init(n);
for(auto e : edges) {
auto path = tree.getPathtoAncestor(e.second, 0);
std::vector<std::pair<int, int>> arr;
for(auto chain : path) {
auto col = colors.upd(chain.first, chain.second, c[e.second]);
std::reverse(col.begin(), col.end());
for(auto elem : col) {
arr.emplace_back(elem.v, elem.r - elem.l);
}
}
long long ans = 0;
for(auto elem : arr) {
//std::cout << "got " << elem.first << " with freq " << elem.second << std::endl;
fw.upd(elem.first, elem.second);
ans += (long long) fw.qry(elem.first - 1) * elem.second;
}
for(auto elem : arr) {
fw.upd(elem.first, -elem.second);
}
std::cout << ans << '\n';
}
}
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
2 ms |
376 KB |
Output is correct |
2 |
Correct |
2 ms |
464 KB |
Output is correct |
3 |
Correct |
2 ms |
544 KB |
Output is correct |
4 |
Correct |
3 ms |
616 KB |
Output is correct |
5 |
Correct |
4 ms |
616 KB |
Output is correct |
6 |
Correct |
4 ms |
616 KB |
Output is correct |
7 |
Correct |
3 ms |
616 KB |
Output is correct |
8 |
Correct |
2 ms |
628 KB |
Output is correct |
9 |
Correct |
3 ms |
656 KB |
Output is correct |
10 |
Correct |
3 ms |
668 KB |
Output is correct |
11 |
Correct |
3 ms |
708 KB |
Output is correct |
12 |
Correct |
3 ms |
720 KB |
Output is correct |
13 |
Correct |
3 ms |
720 KB |
Output is correct |
14 |
Correct |
3 ms |
732 KB |
Output is correct |
15 |
Correct |
5 ms |
744 KB |
Output is correct |
16 |
Correct |
5 ms |
884 KB |
Output is correct |
17 |
Correct |
4 ms |
936 KB |
Output is correct |
18 |
Correct |
4 ms |
936 KB |
Output is correct |
19 |
Correct |
3 ms |
936 KB |
Output is correct |
20 |
Correct |
3 ms |
1040 KB |
Output is correct |
21 |
Correct |
3 ms |
1040 KB |
Output is correct |
22 |
Correct |
4 ms |
1040 KB |
Output is correct |
23 |
Correct |
4 ms |
1144 KB |
Output is correct |
24 |
Correct |
4 ms |
1152 KB |
Output is correct |
25 |
Correct |
5 ms |
1152 KB |
Output is correct |
26 |
Correct |
2 ms |
1152 KB |
Output is correct |
27 |
Correct |
3 ms |
1152 KB |
Output is correct |
28 |
Correct |
3 ms |
1152 KB |
Output is correct |
29 |
Correct |
3 ms |
1152 KB |
Output is correct |
30 |
Correct |
4 ms |
1152 KB |
Output is correct |
31 |
Correct |
4 ms |
1152 KB |
Output is correct |
32 |
Correct |
3 ms |
1152 KB |
Output is correct |
33 |
Correct |
3 ms |
1152 KB |
Output is correct |
34 |
Correct |
2 ms |
1152 KB |
Output is correct |
35 |
Correct |
3 ms |
1168 KB |
Output is correct |
36 |
Correct |
3 ms |
1168 KB |
Output is correct |
37 |
Correct |
3 ms |
1168 KB |
Output is correct |
38 |
Correct |
3 ms |
1180 KB |
Output is correct |
39 |
Correct |
3 ms |
1184 KB |
Output is correct |
40 |
Correct |
3 ms |
1184 KB |
Output is correct |
41 |
Correct |
3 ms |
1184 KB |
Output is correct |
42 |
Correct |
3 ms |
1184 KB |
Output is correct |
43 |
Correct |
3 ms |
1184 KB |
Output is correct |
44 |
Correct |
3 ms |
1184 KB |
Output is correct |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
2 ms |
376 KB |
Output is correct |
2 |
Correct |
2 ms |
464 KB |
Output is correct |
3 |
Correct |
2 ms |
544 KB |
Output is correct |
4 |
Correct |
3 ms |
616 KB |
Output is correct |
5 |
Correct |
4 ms |
616 KB |
Output is correct |
6 |
Correct |
4 ms |
616 KB |
Output is correct |
7 |
Correct |
3 ms |
616 KB |
Output is correct |
8 |
Correct |
2 ms |
628 KB |
Output is correct |
9 |
Correct |
3 ms |
656 KB |
Output is correct |
10 |
Correct |
3 ms |
668 KB |
Output is correct |
11 |
Correct |
3 ms |
708 KB |
Output is correct |
12 |
Correct |
3 ms |
720 KB |
Output is correct |
13 |
Correct |
3 ms |
720 KB |
Output is correct |
14 |
Correct |
3 ms |
732 KB |
Output is correct |
15 |
Correct |
5 ms |
744 KB |
Output is correct |
16 |
Correct |
5 ms |
884 KB |
Output is correct |
17 |
Correct |
4 ms |
936 KB |
Output is correct |
18 |
Correct |
4 ms |
936 KB |
Output is correct |
19 |
Correct |
3 ms |
936 KB |
Output is correct |
20 |
Correct |
3 ms |
1040 KB |
Output is correct |
21 |
Correct |
3 ms |
1040 KB |
Output is correct |
22 |
Correct |
4 ms |
1040 KB |
Output is correct |
23 |
Correct |
4 ms |
1144 KB |
Output is correct |
24 |
Correct |
4 ms |
1152 KB |
Output is correct |
25 |
Correct |
5 ms |
1152 KB |
Output is correct |
26 |
Correct |
2 ms |
1152 KB |
Output is correct |
27 |
Correct |
3 ms |
1152 KB |
Output is correct |
28 |
Correct |
3 ms |
1152 KB |
Output is correct |
29 |
Correct |
3 ms |
1152 KB |
Output is correct |
30 |
Correct |
4 ms |
1152 KB |
Output is correct |
31 |
Correct |
4 ms |
1152 KB |
Output is correct |
32 |
Correct |
3 ms |
1152 KB |
Output is correct |
33 |
Correct |
3 ms |
1152 KB |
Output is correct |
34 |
Correct |
2 ms |
1152 KB |
Output is correct |
35 |
Correct |
3 ms |
1168 KB |
Output is correct |
36 |
Correct |
3 ms |
1168 KB |
Output is correct |
37 |
Correct |
3 ms |
1168 KB |
Output is correct |
38 |
Correct |
3 ms |
1180 KB |
Output is correct |
39 |
Correct |
3 ms |
1184 KB |
Output is correct |
40 |
Correct |
3 ms |
1184 KB |
Output is correct |
41 |
Correct |
3 ms |
1184 KB |
Output is correct |
42 |
Correct |
3 ms |
1184 KB |
Output is correct |
43 |
Correct |
3 ms |
1184 KB |
Output is correct |
44 |
Correct |
3 ms |
1184 KB |
Output is correct |
45 |
Correct |
5 ms |
1320 KB |
Output is correct |
46 |
Correct |
18 ms |
1852 KB |
Output is correct |
47 |
Correct |
20 ms |
1928 KB |
Output is correct |
48 |
Correct |
19 ms |
2004 KB |
Output is correct |
49 |
Correct |
6 ms |
2080 KB |
Output is correct |
50 |
Correct |
6 ms |
2156 KB |
Output is correct |
51 |
Correct |
6 ms |
2232 KB |
Output is correct |
52 |
Correct |
9 ms |
2436 KB |
Output is correct |
53 |
Correct |
9 ms |
2492 KB |
Output is correct |
54 |
Correct |
9 ms |
2568 KB |
Output is correct |
55 |
Correct |
9 ms |
2644 KB |
Output is correct |
56 |
Correct |
8 ms |
2644 KB |
Output is correct |
57 |
Correct |
30 ms |
2692 KB |
Output is correct |
58 |
Correct |
35 ms |
2724 KB |
Output is correct |
59 |
Correct |
30 ms |
2796 KB |
Output is correct |
60 |
Correct |
30 ms |
2876 KB |
Output is correct |
61 |
Correct |
9 ms |
3080 KB |
Output is correct |
62 |
Correct |
9 ms |
3080 KB |
Output is correct |
63 |
Correct |
9 ms |
3112 KB |
Output is correct |
64 |
Correct |
17 ms |
3112 KB |
Output is correct |
65 |
Correct |
18 ms |
3112 KB |
Output is correct |
66 |
Correct |
18 ms |
3152 KB |
Output is correct |
67 |
Correct |
18 ms |
3200 KB |
Output is correct |
68 |
Correct |
6 ms |
3260 KB |
Output is correct |
69 |
Correct |
9 ms |
3564 KB |
Output is correct |
70 |
Correct |
9 ms |
3752 KB |
Output is correct |
71 |
Correct |
8 ms |
3752 KB |
Output is correct |
72 |
Correct |
29 ms |
3752 KB |
Output is correct |
73 |
Correct |
33 ms |
3752 KB |
Output is correct |
74 |
Correct |
8 ms |
3752 KB |
Output is correct |
75 |
Correct |
9 ms |
3872 KB |
Output is correct |
76 |
Correct |
9 ms |
3872 KB |
Output is correct |
77 |
Correct |
9 ms |
3872 KB |
Output is correct |
78 |
Correct |
8 ms |
3968 KB |
Output is correct |
79 |
Correct |
9 ms |
3968 KB |
Output is correct |
80 |
Correct |
13 ms |
3968 KB |
Output is correct |
81 |
Correct |
10 ms |
4112 KB |
Output is correct |
82 |
Correct |
11 ms |
4176 KB |
Output is correct |
83 |
Correct |
11 ms |
4176 KB |
Output is correct |
84 |
Correct |
13 ms |
4304 KB |
Output is correct |
85 |
Correct |
10 ms |
4304 KB |
Output is correct |
86 |
Correct |
10 ms |
4304 KB |
Output is correct |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
2 ms |
376 KB |
Output is correct |
2 |
Correct |
2 ms |
464 KB |
Output is correct |
3 |
Correct |
2 ms |
544 KB |
Output is correct |
4 |
Correct |
3 ms |
616 KB |
Output is correct |
5 |
Correct |
4 ms |
616 KB |
Output is correct |
6 |
Correct |
4 ms |
616 KB |
Output is correct |
7 |
Correct |
3 ms |
616 KB |
Output is correct |
8 |
Correct |
2 ms |
628 KB |
Output is correct |
9 |
Correct |
3 ms |
656 KB |
Output is correct |
10 |
Correct |
3 ms |
668 KB |
Output is correct |
11 |
Correct |
3 ms |
708 KB |
Output is correct |
12 |
Correct |
3 ms |
720 KB |
Output is correct |
13 |
Correct |
3 ms |
720 KB |
Output is correct |
14 |
Correct |
3 ms |
732 KB |
Output is correct |
15 |
Correct |
5 ms |
744 KB |
Output is correct |
16 |
Correct |
5 ms |
884 KB |
Output is correct |
17 |
Correct |
4 ms |
936 KB |
Output is correct |
18 |
Correct |
4 ms |
936 KB |
Output is correct |
19 |
Correct |
3 ms |
936 KB |
Output is correct |
20 |
Correct |
3 ms |
1040 KB |
Output is correct |
21 |
Correct |
3 ms |
1040 KB |
Output is correct |
22 |
Correct |
4 ms |
1040 KB |
Output is correct |
23 |
Correct |
4 ms |
1144 KB |
Output is correct |
24 |
Correct |
4 ms |
1152 KB |
Output is correct |
25 |
Correct |
5 ms |
1152 KB |
Output is correct |
26 |
Correct |
2 ms |
1152 KB |
Output is correct |
27 |
Correct |
3 ms |
1152 KB |
Output is correct |
28 |
Correct |
3 ms |
1152 KB |
Output is correct |
29 |
Correct |
3 ms |
1152 KB |
Output is correct |
30 |
Correct |
4 ms |
1152 KB |
Output is correct |
31 |
Correct |
4 ms |
1152 KB |
Output is correct |
32 |
Correct |
3 ms |
1152 KB |
Output is correct |
33 |
Correct |
3 ms |
1152 KB |
Output is correct |
34 |
Correct |
2 ms |
1152 KB |
Output is correct |
35 |
Correct |
3 ms |
1168 KB |
Output is correct |
36 |
Correct |
3 ms |
1168 KB |
Output is correct |
37 |
Correct |
3 ms |
1168 KB |
Output is correct |
38 |
Correct |
3 ms |
1180 KB |
Output is correct |
39 |
Correct |
3 ms |
1184 KB |
Output is correct |
40 |
Correct |
3 ms |
1184 KB |
Output is correct |
41 |
Correct |
3 ms |
1184 KB |
Output is correct |
42 |
Correct |
3 ms |
1184 KB |
Output is correct |
43 |
Correct |
3 ms |
1184 KB |
Output is correct |
44 |
Correct |
3 ms |
1184 KB |
Output is correct |
45 |
Correct |
5 ms |
1320 KB |
Output is correct |
46 |
Correct |
18 ms |
1852 KB |
Output is correct |
47 |
Correct |
20 ms |
1928 KB |
Output is correct |
48 |
Correct |
19 ms |
2004 KB |
Output is correct |
49 |
Correct |
6 ms |
2080 KB |
Output is correct |
50 |
Correct |
6 ms |
2156 KB |
Output is correct |
51 |
Correct |
6 ms |
2232 KB |
Output is correct |
52 |
Correct |
9 ms |
2436 KB |
Output is correct |
53 |
Correct |
9 ms |
2492 KB |
Output is correct |
54 |
Correct |
9 ms |
2568 KB |
Output is correct |
55 |
Correct |
9 ms |
2644 KB |
Output is correct |
56 |
Correct |
8 ms |
2644 KB |
Output is correct |
57 |
Correct |
30 ms |
2692 KB |
Output is correct |
58 |
Correct |
35 ms |
2724 KB |
Output is correct |
59 |
Correct |
30 ms |
2796 KB |
Output is correct |
60 |
Correct |
30 ms |
2876 KB |
Output is correct |
61 |
Correct |
9 ms |
3080 KB |
Output is correct |
62 |
Correct |
9 ms |
3080 KB |
Output is correct |
63 |
Correct |
9 ms |
3112 KB |
Output is correct |
64 |
Correct |
17 ms |
3112 KB |
Output is correct |
65 |
Correct |
18 ms |
3112 KB |
Output is correct |
66 |
Correct |
18 ms |
3152 KB |
Output is correct |
67 |
Correct |
18 ms |
3200 KB |
Output is correct |
68 |
Correct |
6 ms |
3260 KB |
Output is correct |
69 |
Correct |
9 ms |
3564 KB |
Output is correct |
70 |
Correct |
9 ms |
3752 KB |
Output is correct |
71 |
Correct |
8 ms |
3752 KB |
Output is correct |
72 |
Correct |
29 ms |
3752 KB |
Output is correct |
73 |
Correct |
33 ms |
3752 KB |
Output is correct |
74 |
Correct |
8 ms |
3752 KB |
Output is correct |
75 |
Correct |
9 ms |
3872 KB |
Output is correct |
76 |
Correct |
9 ms |
3872 KB |
Output is correct |
77 |
Correct |
9 ms |
3872 KB |
Output is correct |
78 |
Correct |
8 ms |
3968 KB |
Output is correct |
79 |
Correct |
9 ms |
3968 KB |
Output is correct |
80 |
Correct |
13 ms |
3968 KB |
Output is correct |
81 |
Correct |
10 ms |
4112 KB |
Output is correct |
82 |
Correct |
11 ms |
4176 KB |
Output is correct |
83 |
Correct |
11 ms |
4176 KB |
Output is correct |
84 |
Correct |
13 ms |
4304 KB |
Output is correct |
85 |
Correct |
10 ms |
4304 KB |
Output is correct |
86 |
Correct |
10 ms |
4304 KB |
Output is correct |
87 |
Correct |
51 ms |
5344 KB |
Output is correct |
88 |
Correct |
220 ms |
8860 KB |
Output is correct |
89 |
Correct |
789 ms |
20992 KB |
Output is correct |
90 |
Correct |
805 ms |
23120 KB |
Output is correct |
91 |
Correct |
846 ms |
25264 KB |
Output is correct |
92 |
Correct |
145 ms |
27616 KB |
Output is correct |
93 |
Correct |
140 ms |
29804 KB |
Output is correct |
94 |
Correct |
143 ms |
31992 KB |
Output is correct |
95 |
Correct |
284 ms |
39228 KB |
Output is correct |
96 |
Correct |
294 ms |
41540 KB |
Output is correct |
97 |
Correct |
304 ms |
43584 KB |
Output is correct |
98 |
Correct |
301 ms |
45568 KB |
Output is correct |
99 |
Correct |
275 ms |
45568 KB |
Output is correct |
100 |
Correct |
1315 ms |
45568 KB |
Output is correct |
101 |
Correct |
1296 ms |
45740 KB |
Output is correct |
102 |
Correct |
1317 ms |
47920 KB |
Output is correct |
103 |
Correct |
1321 ms |
50132 KB |
Output is correct |
104 |
Correct |
228 ms |
53664 KB |
Output is correct |
105 |
Correct |
237 ms |
55668 KB |
Output is correct |
106 |
Correct |
230 ms |
57712 KB |
Output is correct |
107 |
Correct |
713 ms |
57712 KB |
Output is correct |
108 |
Correct |
745 ms |
59076 KB |
Output is correct |
109 |
Correct |
801 ms |
60888 KB |
Output is correct |
110 |
Correct |
163 ms |
62612 KB |
Output is correct |
111 |
Correct |
289 ms |
69984 KB |
Output is correct |
112 |
Correct |
292 ms |
71468 KB |
Output is correct |
113 |
Correct |
274 ms |
71468 KB |
Output is correct |
114 |
Correct |
1306 ms |
71468 KB |
Output is correct |
115 |
Correct |
1347 ms |
71468 KB |
Output is correct |
116 |
Correct |
228 ms |
73328 KB |
Output is correct |
117 |
Correct |
305 ms |
74024 KB |
Output is correct |
118 |
Correct |
309 ms |
75484 KB |
Output is correct |
119 |
Correct |
306 ms |
77168 KB |
Output is correct |
120 |
Correct |
297 ms |
79140 KB |
Output is correct |
121 |
Correct |
276 ms |
80260 KB |
Output is correct |
122 |
Correct |
286 ms |
81804 KB |
Output is correct |
123 |
Correct |
259 ms |
84160 KB |
Output is correct |
124 |
Correct |
268 ms |
85564 KB |
Output is correct |
125 |
Correct |
275 ms |
87196 KB |
Output is correct |
126 |
Correct |
250 ms |
89296 KB |
Output is correct |
127 |
Correct |
283 ms |
90472 KB |
Output is correct |
128 |
Correct |
277 ms |
91796 KB |
Output is correct |