#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/pb_ds/detail/standard_policies.hpp>
using namespace std;
using namespace __gnu_pbds;
mt19937_64 rng((int)std::chrono::steady_clock::now().time_since_epoch().count());
const int MAXN = 1e5 + 10;
const int MOD = 1e9 + 7;
#define int long long
typedef tree<pair<int, int>, null_type, less<pair<int, int> >, rb_tree_tag, tree_order_statistics_node_update> ordered_set;
int rnd(int x, int y) {
int u = uniform_int_distribution<int>(x, y)(rng); return u;
}
int n, a[MAXN]; vector<pair<int, int> > adj[MAXN];
bool glo[MAXN];
//Centroid: when you cut this vertex (remove it and
//remove all edges from this vertex), the size of the largest
//connected component of the remaining graph is the smallest possible.
//https://codeforces.com/problemset/problem/1406/C
//Total time complexity: O(n log^2 n)
int sz, mi, ctd;
int ans = 0; // ctd = centroid
int dfs0(int node, int prv) {
int tot = 1;
for(auto x: adj[node]) {
if(!glo[x.first] && x.first != prv) {
tot += dfs0(x.first, node);
}
}
return tot;
}
int dfs1(int node, int prv) {
int ret = 1;
int ma = 0;
for(auto x: adj[node]) {
if(!glo[x.first] && x.first != prv) {
int q = dfs1(x.first, node);
ret += q;
ma = max(ma, q);
}
}
int rem = sz - ret;
ma = max(ma, rem);
if(ma < mi) {
mi = ma;
ctd = node;
}
return ret;
}
int id = 0;
void decomp(int st, int owo) {
sz = owo;
//st: an arbitrary node in the tree in question
//step 1: find the centroid + count number of nodes
mi = 1e9, ctd = -1;
dfs1(st, -1);
//step 2: if nodes = 1, break
if(sz == 1) return;
//step 3: root the subtree at centroid, do a bfs
// all weights = {sum, min}
//step 3a: count id[u] < id[v]
vector<pair<int,int> > all;
for(auto x: adj[ctd]) if(!glo[x.first]) all.push_back(x);
ordered_set s;
glo[ctd] = 1;
vector<int> qwq;
for(int i=0; i<all.size(); i++) {
unordered_map<int, pair<int, int> > w1; // weights to count
unordered_map<int, pair<int, int> > w2; // weights to insert
queue<int> q;
unordered_map<int, int> vis;
vector<int> omg;
w1[all[i].first] = {a[ctd] + all[i].second, a[ctd] + all[i].second};
w2[all[i].first] = {a[all[i].first] + all[i].second, a[all[i].first] + all[i].second};
q.push(all[i].first);
while(q.size()) {
int f=q.front(); q.pop();
if(vis[f] || glo[f]) continue;
vis[f] = 1; omg.push_back(f);
if(s.size()>0) {
if(w1[f].second>=0) ans += s.size();
else ans += s.size() - s.order_of_key({-w1[f].second, 0});
}
for(auto x: adj[f]) {
if(vis[x.first] || glo[x.first]) continue;
q.push(x.first);
w1[x.first] = {w1[f].first + a[f] + x.second, min(w1[f].second, w1[f].first + a[f] + x.second)};
w2[x.first] = {w2[f].first + a[x.first] + x.second, min(w2[f].second, 0ll) + a[x.first] + x.second};
}
}
qwq.push_back(omg.size());
//insert
for(int x: omg) {
if(w2[x].second >= 0) s.insert({w2[x].first, ++id});
}
}
reverse(all.begin(), all.end());
reverse(qwq.begin(), qwq.end());
s.clear();
//step 3b: count id[u] > id[v]
unordered_map<int, pair<int, int> > w1; // weights to count
unordered_map<int, pair<int, int> > w2; // weights to insert
unordered_map<int, int> vis;
queue<int> q;
for(int i=0; i<all.size(); i++) {
w1.clear(); w2.clear();
vis.clear();
vector<int> omg;
w1[all[i].first] = {a[ctd] + all[i].second, a[ctd] + all[i].second};
w2[all[i].first] = {a[all[i].first] + all[i].second, a[all[i].first] + all[i].second};
q.push(all[i].first);
while(q.size()) {
int f=q.front(); q.pop();
if(vis[f] || glo[f]) continue;
vis[f] = 1; omg.push_back(f);
if(s.size()>0) {
if(w1[f].second>=0) ans += s.size();
else ans += s.size() - s.order_of_key({-w1[f].second, 0});
}
for(auto x: adj[f]) {
if(vis[x.first] || glo[x.first]) continue;
q.push(x.first);
w1[x.first] = {w1[f].first + a[f] + x.second, min(w1[f].second, w1[f].first + a[f] + x.second)};
w2[x.first] = {w2[f].first + a[x.first] + x.second, min(w2[f].second, 0ll) + a[x.first] + x.second};
}
}
//insert
for(int x: omg) {
if(w2[x].second >= 0) {
s.insert({w2[x].first, ++id});
}
}
//step 3c: count u = root (= ctd)
for(int x: omg) {
if(w1[x].second >= 0) {
ans++;
}
}
}
//step 3d: count v = root (= ctd)
ans += s.size();
//step 4: remove all edges from the centroid
for(int i=0; i<all.size(); i++) {
if(!glo[all[i].first]) decomp(all[i].first, qwq[i]);
}
}
static char pbuf[1<<19],*p1=pbuf,*p2=pbuf,obuf[1<<19],*o=obuf;
#define gc() p1==p2&&(p2=(p1=pbuf)+fread(pbuf,1,1<<19,stdin),p1==p2)?EOF:*p1++
#define pc(x) (o-obuf<1<<19)?(*o++=(x)):(fwrite(obuf,o-obuf,1,stdout),o=obuf,*o++=(x))
inline int read(){
int x=0;register int c=gc();
while(c<48||c>57)c=gc();
while(c>=48&&c<=57){x=(x<<1)+(x<<3)+c-48;c=gc();}
return x;
}
void solve(int tc) {
n = read();//cin >> n;
for(int i=1; i<=n; i++) a[i] = read(); //cin >> a[i];
for(int i=1; i<n; i++) {
int u, v, w;
u = read();
v = read();
w = read();
//cin >> u >> v >> w;
w = -w;
adj[u].push_back({v, w});
adj[v].push_back({u, w});
}
decomp(1, n);
cout << ans << "\n";
}
int32_t main(){
ios::sync_with_stdio(0); cin.tie(0);
int t = 1; //cin >> t;
for(int i=1; i<=t; i++) solve(i);
}
Compilation message
transport.cpp: In function 'void decomp(long long int, long long int)':
transport.cpp:70:17: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
70 | for(int i=0; i<all.size(); i++) {
| ~^~~~~~~~~~~
transport.cpp:110:17: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
110 | for(int i=0; i<all.size(); i++) {
| ~^~~~~~~~~~~
transport.cpp:150:17: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
150 | for(int i=0; i<all.size(); i++) {
| ~^~~~~~~~~~~
transport.cpp: At global scope:
transport.cpp:154:56: warning: 'o' defined but not used [-Wunused-variable]
154 | static char pbuf[1<<19],*p1=pbuf,*p2=pbuf,obuf[1<<19],*o=obuf;
| ^
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
20 ms |
3412 KB |
Output is correct |
| 2 |
Correct |
15 ms |
3412 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
30 ms |
4052 KB |
Output is correct |
| 2 |
Correct |
31 ms |
4376 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
461 ms |
17900 KB |
Output is correct |
| 2 |
Correct |
306 ms |
14772 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
516 ms |
21148 KB |
Output is correct |
| 2 |
Correct |
549 ms |
23288 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
767 ms |
27748 KB |
Output is correct |
| 2 |
Correct |
854 ms |
34160 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
296 ms |
8500 KB |
Output is correct |
| 2 |
Correct |
124 ms |
7432 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
255 ms |
9372 KB |
Output is correct |
| 2 |
Correct |
361 ms |
10164 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
554 ms |
11004 KB |
Output is correct |
| 2 |
Correct |
496 ms |
13212 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
730 ms |
14856 KB |
Output is correct |
| 2 |
Correct |
659 ms |
16196 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1063 ms |
22200 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
