Submission #232848

#	Time	Username	Problem	Language	Result	Execution time	Memory
232848		Vimmer	Transport (COCI19_transport)	C++14	130 / 130	422 ms	12648 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];

int a[N], n, siz[N], kol;

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

bool mk[N];

vector <ll> se;

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

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

void dfs(int v, int 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(int v, int p, ll cur, ll sm)
{
    cur += ll(a[v]);

    se.pb(sm);

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

        rec(it.F, v, max(0ll, cur - ll(it.S)), sm + max(0ll, ll(it.S) - cur));
    }
}

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

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

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

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

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

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

    int 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(int v, int 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) {kol++; ans += ll(sum(upper_bound(se.begin(), se.end(), sm) - se.begin() - 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(int v)
{
    kol = 0;

    mk[v] = 1;

    se.clear();

    se.pb(-1);

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

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

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

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

    for (int i = 0; i < sz(se); i++) t[i] = 0;

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

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

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

    for (int i = 0; i < sz(se); i++) t[i] = 0;

    for (int i = sz(g[v]) - 1; i >= 0; i--)
    {
        auto it = g[v][i];

        if (mk[it.F]) continue;

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

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

    ans += kol / 2;

    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);

        int 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 (int i = 1; i <= n; i++) cin >> a[i];

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

        cin >> x >> y >> z;

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

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

    dfs(1, -1);

    int root = fnd_cent(1, -1);

    calc(root);

    cout << ans << endl;
}

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