#include <bits/stdc++.h>
#define mp make_pair
#define X first
#define Y second
#define taskname "Transport"
using namespace std;
typedef long long ll;
typedef pair <int, int> ii;
const int N = 1e5 + 1;
struct Tbit{
int d[N];
void update(int x, int val){
for(; x < N; x += x & -x)
d[x] += val;
}
int get(int x){
int res = 0;
for(; x > 0; x -= x & -x)
res += d[x];
return res;
}
} bit;
int n, a[N];
vector <ii> adj[N];
bool done[N];
int sz[N];
ll sumA[N], sumW[N], down[N], up[N];
vector <int> vst;
vector <ll> vals;
void readInput(){
scanf("%d", &n);
for(int i = 1; i <= n; i++)
scanf("%d", &a[i]);
for(int i = 1; i < n; i++){
int u, v, w;
scanf("%d %d %d", &u, &v, &w);
adj[u].push_back(mp(v, w));
adj[v].push_back(mp(u, w));
}
}
void dfsForCentroid(int u, int p = 0){
sz[u] = 1;
for(ii e : adj[u]){
int v = e.X;
if (v == p || done[v])
continue;
dfsForCentroid(v, u);
sz[u] += sz[v];
}
}
int findCentroid(int u, int p, int root){
for(ii e : adj[u]){
int v = e.X;
if (v == p || done[v])
continue;
if (sz[v] >= sz[root] / 2)
return findCentroid(v, u, root);
}
return u;
}
void dfs(int u, int p = 0){
vst.push_back(u);
for(ii e : adj[u]){
int v = e.X;
int w = e.Y;
if (v == p || done[v])
continue;
sumA[v] = sumA[u] + a[v];
sumW[v] = sumW[u] + w;
down[v] = max(down[u], sumW[v] - sumA[u]);
up[v] = max((ll) w - a[v], up[u] + w - a[v]);
dfs(v, u);
}
}
void upd(int u, int p, int type){
int ind = lower_bound(vals.begin(), vals.end(), down[u]) - vals.begin() + 1;
bit.update(ind, type);
for(ii e : adj[u]){
int v = e.X;
if (v == p || done[v])
continue;
upd(v, u, type);
}
}
int get(int u, int p = 0){
int res = 0;
if (up[u] <= 0){
ll val = sumA[u] - sumW[u];
int ind = upper_bound(vals.begin(), vals.end(), val) - vals.begin();
res += bit.get(ind);
}
for(ii e : adj[u]){
int v = e.X;
if (v == p || done[v])
continue;
res += get(v, u);
}
return res;
}
ll calc(int u){
dfsForCentroid(u);
u = findCentroid(u, 0, u);
vst.clear();
vals.clear();
sumW[u] = 0;
sumA[u] = a[u];
up[u] = down[u] = 0;
dfs(u);
for(int v : vst){
vals.push_back(down[v]);
sumA[v] -= a[u];
}
sort(vals.begin(), vals.end());
ll res = 0;
for(int trial = 0; trial < 2; trial++){
if (trial == 0){
res += bit.get(1);
bit.update(1, 1);
}
for(ii e : adj[u]){
int v = e.X;
if (done[v])
continue;
res += get(v, u);
upd(v, u, 1);
}
if (trial == 1){
res += bit.get(1);
bit.update(1, 1);
}
for(ii e : adj[u]){
int v = e.X;
if (done[v])
continue;
upd(v, u, -1);
}
bit.update(1, -1);
reverse(adj[u].begin(), adj[u].end());
}
done[u] = 1;
for(ii e : adj[u]){
int v = e.X;
if (!done[v])
res += calc(v);
}
return res;
}
void solve(){
cout << calc(1);
}
int main(){
//freopen(taskname".inp", "r", stdin);
//freopen(taskname".out", "w", stdout);
readInput();
solve();
return 0;
}
Compilation message
transport.cpp: In function 'void readInput()':
transport.cpp:39:10: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
39 | scanf("%d", &n);
| ~~~~~^~~~~~~~~~
transport.cpp:41:14: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
41 | scanf("%d", &a[i]);
| ~~~~~^~~~~~~~~~~~~
transport.cpp:44:14: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
44 | scanf("%d %d %d", &u, &v, &w);
| ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
14 ms |
3072 KB |
Output is correct |
| 2 |
Correct |
12 ms |
3200 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
18 ms |
3456 KB |
Output is correct |
| 2 |
Correct |
22 ms |
3584 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
221 ms |
9212 KB |
Output is correct |
| 2 |
Correct |
174 ms |
8716 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
282 ms |
11696 KB |
Output is correct |
| 2 |
Correct |
310 ms |
13044 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
408 ms |
14960 KB |
Output is correct |
| 2 |
Correct |
490 ms |
17644 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
158 ms |
5240 KB |
Output is correct |
| 2 |
Correct |
81 ms |
5120 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
148 ms |
7280 KB |
Output is correct |
| 2 |
Correct |
284 ms |
6644 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
418 ms |
8436 KB |
Output is correct |
| 2 |
Correct |
393 ms |
8872 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
602 ms |
10100 KB |
Output is correct |
| 2 |
Correct |
599 ms |
10348 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
715 ms |
12396 KB |
Output is correct |
| 2 |
Correct |
642 ms |
11764 KB |
Output is correct |
