Submission #232835

#	Time	Username	Problem	Language	Result	Execution time	Memory
232835		Vimmer	Transport (COCI19_transport)	C++14	0 / 130	9 ms	3584 KiB

Main

#include <bits/stdc++.h>
//#include <ext/pb_ds/assoc_container.hpp>
//#include <ext/pb_ds/tree_policy.hpp>

//#pragma GCC optimize("unroll-loops")
//#pragma GCC optimize("-O3")
//#pragma GCC optimize("Ofast")

#define sz(x) ll(x.size())
#define N 100015
#define base 1000000
#define M ll(1e9+7)
#define F first
#define S second
#define pb push_back
#define in insert
#define eb emplace_back
#define ed "\n"

using namespace std;
//using namespace __gnu_pbds;

typedef long long ll;
typedef long double ld;
typedef short int si;

//typedef tree<ll, null_type, less<ll>, rb_tree_tag, tree_order_statistics_node_update> ordered_set;


ll ans, t[N];

ll a[N], n, siz[N];

vector <pair <ll, ll> > g[N];

bool mk[N];

vector <ll> se;

void add(ll v, ll val) {for (; v < N; v = (v | (v + 1))) t[v] += val;}

void del(ll v, ll val) {for (; v < N; v = (v | (v + 1))) t[v] -= val;}

ll sum(ll v) {ll res = 0; for (; v >= 0; v = (v & (v + 1)) - 1) res += t[v]; return res;}

void dfs(ll v, ll p)
{
    siz[v] = 1;

    for (auto it : g[v])
    {
        if (it.F == p || mk[it.F]) continue;

        dfs(it.F, v);

        siz[v] += siz[it.F];
    }
}

void rec_add(ll v, ll p, ll cur, ll sm, bool f)
{
    cur += ll(a[v]);

    if (f) se.pb(sm); else add(upper_bound(se.begin(), se.end(), sm) - se.begin() - 1, 1);

    for (auto it : g[v])
    {
        if (it.F == p || mk[it.F]) continue;

        rec_add(it.F, v, cur - ll(it.S), max(sm, ll(it.S) - cur), f);
    }
}

void rec_del(ll v, ll p, ll cur, ll sm)
{
    cur += ll(a[v]);

    del(upper_bound(se.begin(), se.end(), sm) - se.begin() - 1, 1);

    for (auto it : g[v])
    {
        if (it.F == p || mk[it.F]) continue;

        rec_del(it.F, v, cur - ll(it.S), max(sm, ll(it.S) - cur));
    }
}

ll fnd_cent(ll v, ll p)
{
    if (p != -1)
    {
        siz[p] -= siz[v];

        siz[v] += siz[p];
    }

    ll nm = -1;

    for (auto it : g[v])
    {
        if (mk[it.F]) continue;

        if (siz[it.F] > siz[v] / 2) {nm = it.F; break;}
    }

    if (nm == -1) return v;

    return fnd_cent(nm, v);
}

void spusk(ll v, ll p, ll sm, ll cost, ll need)
{
    siz[v] = 1;

    need += cost;

    need = max(0ll, need - a[v]);

    sm += a[v];

    sm -= cost;

    if (need == 0) ans += ll(sum(upper_bound(se.begin(), se.end(), sm) - se.begin() - 1)) + 1;

    for (auto it : g[v])
    {
        if (it.F == p || mk[it.F]) continue;

        spusk(it.F, v, sm, it.S, need);

        siz[v] += siz[it.F];
    }
}

void calc(ll v)
{
    mk[v] = 1;

    se.clear();

    memset(t, 0, sizeof(t));

    se.pb(-1);

    for (auto it : g[v])
    {
        if (mk[it.F]) continue;

        rec_add(it.F, -1, 0, it.S, 1);
    }

    sort(se.begin(), se.end());

    se.resize(unique(se.begin(), se.end()) - se.begin());

    for (auto it : g[v])
    {
        if (mk[it.F]) continue;

        rec_add(it.F, -1, 0, it.S, 0);
    }

    for (auto it : g[v])
    {
        if (mk[it.F]) continue;

        rec_del(it.F, -1, 0, it.S);

        spusk(it.F, -1, a[v], it.S, 0);

        rec_add(it.F, -1, 0, it.S, 0);
    }

    ans += ll(sum(upper_bound(se.begin(), se.end(), a[v]) - se.begin() - 1));

    for (auto it : g[v])
    {
        if (mk[it.F]) continue;

        dfs(it.F, -1);

        ll root = fnd_cent(it.F, -1);

        if (root != -1) calc(root);
    }
}

int main()
{
    freopen("input.txt","r",stdin);// freopen("output.txt","w",stdout);

    ios_base::sync_with_stdio(0); istream::sync_with_stdio(0); cin.tie(0); cout.tie(0);

    cin >> n;

    for (ll i = 1; i <= n; i++) cin >> a[i];

    for (ll i = 1; i < n; i++)
    {
        ll x, y, z;

        cin >> x >> y >> z;

        g[x].pb({y, z});

        g[y].pb({x, z});
    }

    dfs(1, -1);

    ll root = fnd_cent(1, -1);

    calc(root);

    cout << ans << endl;
}

Compilation message (stderr)

transport.cpp: In function 'int main()':
transport.cpp:190:12: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)', declared with attribute warn_unused_result [-Wunused-result]
     freopen("input.txt","r",stdin);// freopen("output.txt","w",stdout);
     ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~

• Subtask 1: 0 / 13
• Subtask 2: 0 / 13
• Subtask 3: 0 / 13
• Subtask 4: 0 / 13
• Subtask 5: 0 / 13
• Subtask 6: 0 / 13
• Subtask 7: 0 / 13
• Subtask 8: 0 / 13
• Subtask 9: 0 / 13
• Subtask 10: 0 / 13