#include <bits/stdc++.h>
using namespace std;
const int N = 100005;
int n, poi, cnt, curCentroid;
int subSizes[N];
vector<array<long long, 2>> adj[N];
bool vis[N];
vector<long long> all, v[N], c[N], g[N];
long long ans = 0;
long long a[N];
long long fenw[2 * N];
int times[2 * N];
void modify(int x, long long add) {
while (x < (int) all.size()) {
if (times[x] != cnt) {
times[x] = cnt;
fenw[x] = 0;
}
fenw[x] += add;
x |= (x + 1);
}
}
long long get(int x) {
long long sum = 0;
while (x >= 0) {
if (times[x] != cnt) {
fenw[x] = 0;
}
sum += fenw[x];
x = (x & (x + 1)) - 1;
}
return sum;
}
void initArrays() {
for (int i = 0; i < n; i++) {
vis[i] = false;
}
for (int i = 0; i < 2 * n; i++) {
times[i] = -1;
}
return ;
}
int getSiz(int u, int prev) {
subSizes[u] = 1;
for (auto x : adj[u]) {
if (!vis[x[0]] && x[0] != prev) {
subSizes[u] += getSiz(x[0], u);
}
}
return subSizes[u];
}
int locateCentroid(int u, int prev, int sz) {
for (auto x : adj[u]) {
if (vis[x[0]] || x[0] == prev) {
continue;
}
if (subSizes[x[0]] > sz / 2) {
return locateCentroid(x[0], u, sz);
}
}
return u;
}
int centroid(int u, int prev) {
int sz = getSiz(u, prev);
return locateCentroid(u, prev, sz);
}
void dfs(int u, int prev, long long balance, long long minPrefixBalance, long long minSuffixBalance) {
all.push_back(-minPrefixBalance + a[curCentroid]);
v[poi].push_back(minPrefixBalance);
minSuffixBalance += a[u];
balance += a[u];
all.push_back(balance);
g[poi].push_back(minSuffixBalance);
c[poi].push_back(balance);
for (auto x : adj[u]) {
if (x[0] == prev || vis[x[0]]) {
continue;
}
balance += x[1];
long long nMnPref = min(minPrefixBalance, balance);
dfs(x[0], u, balance, nMnPref, min(0LL, minSuffixBalance) + x[1]);
balance -= x[1];
}
return ;
}
void Decompose(int u, int prev) {
u = centroid(u, prev);
curCentroid = u;
vis[u] = 1;
poi = 0;
cnt++;
all.clear();
for (auto x : adj[u]) {
if (x[0] == prev || vis[x[0]]) {
continue;
}
dfs(x[0], u, a[u] + x[1], a[u] + x[1], x[1]);
++poi;
}
// all.push_back(a[u]);
sort(all.begin(), all.end());
all.erase(unique(all.begin(), all.end()), all.end());
auto Compress = [&](long long pts) {
int x = lower_bound(all.begin(), all.end(), pts) - all.begin();
return x;
};
for (int i = 0; i < poi; i++) {
for (int j = 0; j < (int) v[i].size(); j++) {
long long minPrefixBalance = v[i][j] - a[u];
if (minPrefixBalance + a[u] >= 0) {
ans += 1;
}
if (minPrefixBalance < 0) {
modify(Compress(-minPrefixBalance), +1);
} else {
modify(0, +1);
}
}
}
for (int i = 0; i < poi; i++) {
for (int j = 0; j < (int) c[i].size(); j++) {
long long minPrefixBalance = v[i][j] - a[u];
if (minPrefixBalance < 0) {
modify(Compress(-minPrefixBalance), -1);
} else {
modify(0, -1);
}
}
for (int j = 0; j < (int) c[i].size(); j++) {
long long curBalance = c[i][j];
long long minSuffixBalance = g[i][j];
if (minSuffixBalance >= 0) {
ans += 1;
ans += get(Compress(curBalance));
}
}
for (int j = 0; j < (int) c[i].size(); j++) {
long long minPrefixBalance = v[i][j] - a[u];
if (minPrefixBalance < 0) {
modify(Compress(-minPrefixBalance), +1);
} else {
modify(0, +1);
}
}
}
for (int i = 0; i < poi; i++) {
v[i].clear();
c[i].clear();
g[i].clear();
}
for (auto x : adj[u]) {
if (!vis[x[0]] && x[0] != prev) {
Decompose(x[0], u);
}
}
}
int main() {
ios::sync_with_stdio(false);
cin.tie(0);
cin >> n;
initArrays();
for (int i = 0; i < n; i++) {
cin >> a[i];
}
for (int i = 0; i < n - 1; i++) {
int f, t, cost;
cin >> f >> t >> cost;
--f, --t;
adj[f].push_back({t, -cost});
adj[t].push_back({f, -cost});
}
Decompose(0, -1);
cout << ans << '\n';
return 0;
}
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
11 ms |
10196 KB |
Output is correct |
| 2 |
Correct |
12 ms |
10512 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
14 ms |
10580 KB |
Output is correct |
| 2 |
Correct |
17 ms |
10836 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
137 ms |
18964 KB |
Output is correct |
| 2 |
Correct |
127 ms |
18008 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
185 ms |
22224 KB |
Output is correct |
| 2 |
Correct |
206 ms |
24172 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
274 ms |
27360 KB |
Output is correct |
| 2 |
Correct |
315 ms |
30500 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
85 ms |
14196 KB |
Output is correct |
| 2 |
Correct |
55 ms |
14404 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
99 ms |
18452 KB |
Output is correct |
| 2 |
Correct |
133 ms |
16640 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
217 ms |
19720 KB |
Output is correct |
| 2 |
Correct |
194 ms |
20248 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
262 ms |
22968 KB |
Output is correct |
| 2 |
Correct |
295 ms |
23648 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
419 ms |
26640 KB |
Output is correct |
| 2 |
Correct |
366 ms |
24668 KB |
Output is correct |
