#include <bits/stdc++.h>
#define long long long
using namespace std;
const int N = 1e5+1;
void update(int i, vector<long> &fen) {
while(i < fen.size()) ++fen[i], i += i & -i;
}
long query(int i, vector<long> &fen) {
long ret = 0;
while(i > 0) ret += fen[i], i -= i & -i;
return ret;
}
int n, a[N];
vector<pair<int,int> > g[N];
long ans;
int cen, sz[N];
bool blocked[N];
void find_centroid(int u, int par, int tot, pair<int,int> &tmp) {
sz[u] = 1;
int mx = 0;
for(auto p: g[u]) {
int v, w;
tie(v, w) = p;
if(v == par || blocked[v]) continue;
find_centroid(v, u, tot, tmp);
sz[u] += sz[v];
mx = max(sz[v], mx);
}
mx = max(tot - sz[u], mx);
if(mx < tmp.first) tmp = {mx, u};
}
long fup[N], fdown[N], cup[N];
int tick, st[N], en[N];
vector<pair<long, int> > ls;
void dfs1(int u, int par, int child_subtree, long cdown = 0) {
// cout << u << ' ' << child_subtree << endl;
st[u] = ++tick;
for(auto p: g[u]) {
int v, w;
tie(v, w) = p;
if(v == par || blocked[v]) continue;
fup[v] = min(fup[u], 0ll) + a[v] - w;
cup[v] = cup[u] + a[v] - w;
fdown[v] = min(cdown + a[u] - w, fdown[u]);
dfs1(v, u, (st[u] == 1 ? v : child_subtree), cdown + a[u] - w);
}
en[u] = tick;
if(st[u] != 1) {
if(fup[u] >= 0) {
++ans; // path towards centroid = root
ls.emplace_back(-cup[u], -u);
}
if(fdown[u] >= 0) {
++ans;
}
assert(child_subtree != -1);
ls.emplace_back(fdown[u], child_subtree); // be careful of including the centroid!!!
}
// cout << u << ' ' << cup[u] << ' ' << fup[u] << ' ' << fdown[u] << endl;
}
void dec(int src, int tot) {
pair<int,int> tmp = {tot, -1};
find_centroid(src, -1, tot, tmp);
cen = tmp.second;
// cout << "cen " << cen << ' ' << tot << endl;
tick = 0;
ls.clear();
fup[cen] = fdown[cen] = cup[cen] = 0;
dfs1(cen, -1, -1);
sort(ls.begin(), ls.end());
vector<long> fen(tick+1);
long tmp_ans = 0;
for(auto p: ls) {
int v;
tie(ignore, v) = p;
// cout << "ls " << val << ' ' << s << ' ' << e << endl;
if(v < 0) {
update(st[-v], fen);
} else { // don't consider centroid for this section
tmp_ans += query(st[v]-1, fen) + query(tick, fen) - query(en[v], fen);
// cout << "res " << query(s-1, fen) << ' ' <<query(tick, fen) - query(e, fen) << ' ' << query(s-1, fen) + query(tick, fen) - query(e, fen) << endl;
}
}
ans += tmp_ans;
// cout << "curr_ans " << ans << ' ' << tmp_ans << endl;
// cout << query(tick, fen) << endl;
blocked[cen] = true;
for(auto p: g[cen]) {
int v, w;
tie(v, w) = p;
if(blocked[v]) continue;
dec(v, (sz[v] > sz[cen] ? tot - sz[cen] : sz[v]));
}
}
void solve(istream& cin) {
cin >> n;
for(int i = 1; i <= n; i++) cin >> a[i];
for(int i = 1, u, v, w; i < n; i++) {
cin >> u >> v >> w;
g[u].emplace_back(v, w);
g[v].emplace_back(u, w);
}
dec(1, n);
cout << ans << endl;
}
int main() {
ios::sync_with_stdio(false);
cin.tie(0);
// ifstream in("3.in");
solve(cin);
}
Compilation message
transport.cpp: In function 'void update(int, std::vector<long long int>&)':
transport.cpp:7:13: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
while(i < fen.size()) ++fen[i], i += i & -i;
~~^~~~~~~~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
11 ms |
3200 KB |
Output is correct |
| 2 |
Correct |
13 ms |
3328 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
14 ms |
3584 KB |
Output is correct |
| 2 |
Correct |
18 ms |
3780 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
173 ms |
11204 KB |
Output is correct |
| 2 |
Correct |
101 ms |
11108 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
161 ms |
14376 KB |
Output is correct |
| 2 |
Correct |
162 ms |
16692 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
240 ms |
18540 KB |
Output is correct |
| 2 |
Correct |
263 ms |
22652 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
136 ms |
6352 KB |
Output is correct |
| 2 |
Correct |
62 ms |
6108 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
118 ms |
8444 KB |
Output is correct |
| 2 |
Correct |
212 ms |
8768 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
278 ms |
9708 KB |
Output is correct |
| 2 |
Correct |
230 ms |
11116 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
424 ms |
11668 KB |
Output is correct |
| 2 |
Correct |
406 ms |
12904 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
688 ms |
16424 KB |
Output is correct |
| 2 |
Correct |
517 ms |
15100 KB |
Output is correct |
