oj.uz

Submission #943297

# Submission time Handle Problem Language Result Execution time Memory
943297 2024-03-11T10:14:09 Z 12345678 Transport (COCI19_transport) C++17
130 / 130
 521 ms 21444 KB 
#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=0, dp[nx], t, lft[nx];
vector<pair<ll, ll>> d[nx];
map<ll, ll> mp;
 
struct fenwick
{
    ll d[nx];
    void add(ll i, ll x)
    {
        while (i<nx) d[i]+=x, 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, ll pre, bool mode)
{
    if (sm>=0&&pre>=0)
    {
        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, sm+vl[v]-w, min(0ll, pre+(vl[v]-w)), mode);
}
 
void dfsadd(ll u, ll p, bool mode, ll add)
{
    //if (mode&&add==1) cout<<"here "<<u<<' '<<dp[u]<<'\n';
    if (mode) f.add(mp[dp[u]], add);
    else mp[dp[u]]=0;
    for (auto [v, w]:d[u]) 
    {
        if (v!=p&&!used[v]) 
        {
            if (w>lft[u]) dp[v]=dp[u]+(w-lft[u]), lft[v]=0;
            else dp[v]=dp[u], lft[v]=lft[u]-w;
            lft[v]+=vl[v];
            dfsadd(v, u, mode, add);
        }
    }
}
 
void decomposition(ll u)
{
    u=findcentroid(u, u, dfssz(u, u));
    used[u]=1;
    dp[u]=lft[u]=0;
    mp.clear();
    mp[vl[u]]=0;
    t=0;
    //cout<<"root "<<u<<'\n';
    for (auto [v, w]:d[u]) if (!used[v]) dp[v]=w, lft[v]=vl[v], dfsadd(v, u, 0, 0), dfsquery(v, u, u, -w+vl[v], 0, 0);
    for (auto &[x, y]:mp) y=++t;
    f.add(mp[vl[u]], 1);
    for (auto [v, w]:d[u]) if (!used[v]) dfsquery(v, u, u, -w+vl[v], 0, 1), dfsadd(v, u, 1, 1);
    f.add(mp[vl[u]], -1);
    for (auto [v, w]:d[u]) if (!used[v]) dfsadd(v, u, 1, -1);
    reverse(d[u].begin(), d[u].end());
    for (auto [v, w]:d[u]) if (!used[v]) dfsquery(v, u, u, -w+vl[v], 0, 1), dfsadd(v, u, 1, 1);
    ans+=f.query(mp[vl[u]]);
    for (auto [v, w]:d[u]) if (!used[v]) dfsadd(v, u, 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;
}

/*
6
1 2 3 3 2 1
1 2 4
2 3 4
3 4 1
4 5 4
5 6 4
*/

Subtask #1
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 10 ms 6748 KB Output is correct
2 Correct 9 ms 7132 KB Output is correct

Subtask #2
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 14 ms 7256 KB Output is correct
2 Correct 16 ms 7240 KB Output is correct

Subtask #3
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 186 ms 13804 KB Output is correct
2 Correct 126 ms 12280 KB Output is correct

Subtask #4
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 214 ms 15700 KB Output is correct
2 Correct 230 ms 16464 KB Output is correct

Subtask #5
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 318 ms 19464 KB Output is correct
2 Correct 366 ms 21444 KB Output is correct

Subtask #6
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 140 ms 9564 KB Output is correct
2 Correct 70 ms 9308 KB Output is correct

Subtask #7
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 131 ms 12120 KB Output is correct
2 Correct 198 ms 10580 KB Output is correct

Subtask #8
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 292 ms 13136 KB Output is correct
2 Correct 334 ms 13908 KB Output is correct

Subtask #9
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 419 ms 14928 KB Output is correct
2 Correct 424 ms 15424 KB Output is correct

Subtask #10
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 521 ms 16768 KB Output is correct
2 Correct 491 ms 15032 KB Output is correct