This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include <bits/stdc++.h>
using namespace std;
#define ll long long
const ll nx=1e5+5;
ll n, vl[nx], u, v, w, used[nx], sz[nx], ans, t;
queue<ll> q;
vector<pair<ll, ll>> d[nx];
map<ll, ll> mp;
struct fenwick
{
ll d[nx];
void add(ll i, ll vl)
{
while (i<nx) d[i]+=vl, i+=(i&-i);
}
ll query(ll i)
{
ll res=0;
while (i>0) res+=d[i], i-=(i&-i);
return res;
}
} f;
ll dfssz(ll u, ll p)
{
sz[u]=1;
for (auto [v, w]:d[u]) if (v!=p&&!used[v]) sz[u]+=dfssz(v, u);
return sz[u];
}
ll findcentroid(ll u, ll p, ll rtsz)
{
for (auto [v, w]:d[u]) if (v!=p&&!used[v]&&2*sz[v]>rtsz) return findcentroid(v, u, rtsz);
return u;
}
void dfsquery(ll u, ll p, ll rt, ll sm, bool mode)
{
if (sm>=0)
{
//if (mode) cout<<u<<' '<<p<<' '<<sm<<' '<<sm+vl[rt]<<' '<<f.query(mp[sm+vl[rt]])<<'\n';
if (mode) ans+=f.query(mp[sm+vl[rt]]);
else mp[sm+vl[rt]]=0;
}
for (auto [v, w]:d[u]) if (v!=p&&!used[v]) dfsquery(v, u, rt, min(0ll, sm)+vl[v]-w, mode);
}
void dfsadd(ll u, ll p, ll w, ll cost, ll left, bool mode, ll add)
{
//if (u==8&&p==6) cout<<"here "<<w<<' '<<'\n';
if (w>left) cost+=(w-left), left=0;
else left-=w;
if (mode) f.add(mp[cost], add);
else mp[cost]=0;
for (auto [v, w]:d[u]) if (v!=p&&!used[v]) dfsadd(v, u, w, cost, vl[u]+left, mode, add);
}
void decomposition(ll u)
{
u=findcentroid(u, u, dfssz(u, u));
used[u]=1;
mp.clear();
mp[0]=0;
mp[vl[u]]=0;
t=0;
for (auto [v, w]:d[u]) if (!used[v]) dfsadd(v, u, w, 0, 0, 0, 0), dfsquery(v, u, u, -w+vl[v], 0);
for (auto &[x, y]:mp) y=++t;
f.add(mp[0], 1);
for (ll i=0; i<d[u].size(); i++)
{
auto [v, w]=d[u][i];
if (used[v]) continue;
//cout<<"order "<<u<<' '<<v<<'\n';
dfsquery(v, u, u, -w+vl[v], 1);
dfsadd(v, u, w, 0, 0, 1, 1);
}
f.add(mp[0], -1);
for (auto [v, w]:d[u]) if (!used[v]) dfsadd(v, u, w, 0, 0, 1, -1);
for (ll i=d[u].size()-1; i>=0; i--)
{
auto [v, w]=d[u][i];
if (used[v]) continue;
//cout<<"order "<<u<<' '<<v<<'\n';
dfsquery(v, u, u, -w+vl[v], 1);
dfsadd(v, u, w, 0, 0, 1, 1);
}
ans+=f.query(mp[vl[u]]);
//cout<<"ans "<<u<<' '<<ans<<'\n';
for (auto [v, w]:d[u]) if (!used[v]) dfsadd(v, u, w, 0, 0, 1, -1);
for (auto [v, w]:d[u]) if (!used[v]) decomposition(v);
}
int main()
{
cin.tie(NULL)->sync_with_stdio(false);
cin>>n;
for (ll i=1; i<=n; i++) cin>>vl[i];
for (ll i=1; i<n; i++) cin>>u>>v>>w, d[u].push_back({v, w}), d[v].push_back({u, w});
decomposition(1);
cout<<ans;
}
/*
5
3 1 2 4 5
1 2 3
3 2 2
4 2 6
5 4 3
8
5 2 4 7 8 3 3 6
6 5 5
1 4 5
3 1 2
8 6 5
1 2 3
4 5 3
4 7 5
*/
Compilation message (stderr)
transport.cpp: In function 'void decomposition(long long int)':
transport.cpp:74:19: 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]
74 | for (ll i=0; i<d[u].size(); i++)
| ~^~~~~~~~~~~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |