#include <bits/stdc++.h>
using namespace std;
#define f first
#define s second
#define pb push_back
#define ar array
#define all(x) x.begin(), x.end()
#define siz(x) (int)x.size()
#define FOR(x, y, z) for(int x = (y); x < (z); x++)
#define ROF(x, z, y) for(int x = (y-1); x >= (z); x--)
#define F0R(x, z) FOR(x, 0, z)
#define R0F(x, z) ROF(x, 0, z)
#define trav(x, y) for(auto&x:y)
using ll = long long;
using vi = vector<int>;
using vl = vector<long long>;
using pii = pair<int, int>;
using vpii = vector<pair<int, int>>;
template<class T> inline bool ckmin(T&a, T b) {return b < a ? a = b, 1 : 0;}
template<class T> inline bool ckmax(T&a, T b) {return b > a ? a = b, 1 : 0;}
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/pb_ds/assoc_container.hpp>
//#define int long long
using namespace __gnu_pbds;
#define ordered_set tree<pair<ll, ll>, null_type,less<pair<ll, ll>>, rb_tree_tag,tree_order_statistics_node_update>
const char nl = '\n';
const int mxN = 2e5 + 10;
const int MOD = 1e9 + 7;
const long long infLL = 1e18;
struct centroid{
vector<vector<int>> adj;
vector<bool> vis;
vi par, sz;
int n;
void init(int _n){ n = _n; adj.resize(n); par.resize(n); sz.resize(n); vis.resize(n, 0);}
void add(int a, int b) {adj[a].pb(b); adj[b].pb(a);}
int find_size(int u, int p = -1) {
if(vis[u]) return 0;
sz[u] = 1;
trav(v, adj[u]){
if(v != p)
sz[u] += find_size(v, u);
}
return sz[u];
}
int find_centroid(int u, int p, int N){
trav(v, adj[u]){
if(v!=p&&!vis[v]&&sz[v]>N/2)
return find_centroid(v, u, N);
}
return u;
}
void init_centroid(int u = 0, int p = -1){
find_size(u);
int c = find_centroid(u, -1, sz[u]);
vis[c] = 1;
par[c] = p;
trav(v, adj[c])
if(!vis[v])
init_centroid(v, c);
}
};
centroid C;
int n;
vector<pair<int, ll>> adj[mxN];
set<int> padj[mxN];
set<pair<int, int>> al;
set<int> nodes[mxN];
bool active[mxN];
ll cost[mxN];
ll a[mxN];
vl costs; vi to_remove;
ll ans = 0;
void dfs(int u, int p, ll Cost, ll sum){
//cout << u << ' ' << p << endl;
//cout << "cost[u] " << u << ' ' << Cost << ' ' << sum << nl;
cost[u] = Cost;
trav(v, adj[u]){
if(v.f == p || !active[v.f]) continue;
//cout << sum << ' ' << v.f << ' ' << Cost << ' ' << a[u] << nl;
dfs(v.f, u, min(sum+a[u]-v.s, Cost), sum+a[u]-v.s);
}
}
void dfs2(int u, int p, ll Cost, ll sum){
//cout << u << ' ' << Cost << ' ' << sum << nl;
to_remove.pb(u);
if(Cost >= 0){
costs.pb(sum);
}
trav(v, adj[u]){
if(p == v.f || !active[v.f]) continue;
//if(active[v.f] && v.f != u){
dfs2(v.f, u, min(a[v.f] - v.s, Cost + (a[v.f] - v.s)), sum + a[v.f] - v.s);
//}
}
}
int32_t main(){
ios_base::sync_with_stdio(0); cin.tie(0);
cin >> n;
F0R(i, n){
cin >> a[i];
}
C.init(n);
F0R(i, n-1){
int e1, e2; ll w;
cin >> e1 >> e2 >> w;
e1--; e2--;
adj[e1].pb({e2, w});
adj[e2].pb({e1, w});
C.add(e1, e2);
}
C.init_centroid();
F0R(i, n){
if(C.par[i]!=-1)
padj[C.par[i]].insert(i);
}
F0R(i, n){
al.insert({siz(padj[i]), i});
nodes[i].insert(i);
}
while(siz(al)){
pii x = *al.begin();
assert(x.f == 0) ;
al.erase(al.begin());
//cout << x.s << nl;
trav(X, nodes[x.s]){
active[X] = 1;
//cout << "together " << X << ' ' << x.s << nl;
}
cost[x.s] = 0;
//dfs(x.s, -1, 1e18, 0);
trav(c, adj[x.s]){
if(!active[c.f]) continue;
dfs(c.f, x.s, a[x.s]-c.s, a[x.s]-c.s);
}
ordered_set S;
trav(X, nodes[x.s]){
S.insert({cost[X], X});
//cout << "cost " << X << " " << cost[X] << nl;
if(cost[X] >= 0) {
if(X != x.s){
//cout << "WORKS " << X << ' ' << x.s << nl;
}
ans++;
}
}
trav(c, adj[x.s]){
int child = c.f;
if(!active[child]) continue;
to_remove.clear(); costs.clear();
dfs2(child, x.s, a[c.f]-c.s, a[c.f]-c.s);
trav(X, to_remove){
S.erase({cost[X], X});
}
trav(X, costs){
ans += siz(S) - S.order_of_key({-X, -1});
//cout << "stuff " << X << ' ' << siz(S) - S.order_of_key({-X, -1}) << nl;
}
//cout << "S: ";
//trav(ss, S){
// cout << ss.f << ' ';
//}
//cout << nl;
trav(X, to_remove){
S.insert({cost[X], X});
}
}
trav(X, nodes[x.s]){
active[X] = 0;
}
if(C.par[x.s] != -1){
int p = C.par[x.s];
al.erase({siz(padj[p]), p});
padj[p].erase(x.s);
al.insert({siz(padj[p]), p});
if(siz(nodes[p]) < siz(nodes[x.s])){
swap(nodes[p], nodes[x.s]);
}
trav(X, nodes[x.s]){
nodes[p].insert(X);
}
}
}
cout << ans - n << nl;
return 0;
}
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
38 ms |
25324 KB |
Output is correct |
| 2 |
Correct |
36 ms |
25600 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
44 ms |
26092 KB |
Output is correct |
| 2 |
Correct |
54 ms |
26860 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
489 ms |
52204 KB |
Output is correct |
| 2 |
Correct |
412 ms |
50540 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
713 ms |
63796 KB |
Output is correct |
| 2 |
Runtime error |
690 ms |
65540 KB |
Execution killed with signal 9 |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Runtime error |
637 ms |
65540 KB |
Execution killed with signal 9 |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
316 ms |
38460 KB |
Output is correct |
| 2 |
Correct |
199 ms |
34924 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
389 ms |
44268 KB |
Output is correct |
| 2 |
Correct |
481 ms |
47212 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
811 ms |
55916 KB |
Output is correct |
| 2 |
Correct |
741 ms |
58796 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Runtime error |
944 ms |
65540 KB |
Execution killed with signal 9 |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Runtime error |
895 ms |
65536 KB |
Execution killed with signal 9 |
| 2 |
Halted |
0 ms |
0 KB |
- |
