Submission #239660

#	Time	Username	Problem	Language	Result	Execution time	Memory
239660		kartel	Transport (COCI19_transport)	C++14	130 / 130	455 ms	17004 KiB

Main

#include <bits/stdc++.h>
//#include <ext/pb_ds/assoc_container.hpp>
//#include <ext/pb_ds/tree_policy.hpp>
#define in(x) freopen(x, "r", stdin)
#define out(x) freopen(x, "w", stdout)
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
#pragma GCC optimize("-O3")
#define F first
#define S second
#define pb push_back
#define N +100500
#define M ll(1e9 + 7)
#define sz(x) (int)x.size()
#define re return
#define oo ll(1e9)
#define el '\n'
#define pii pair <int, int>
using namespace std;
//using namespace __gnu_pbds;
//typedef tree <int, null_type, less_equal <int> , rb_tree_tag, tree_order_statistics_node_update> ordered_set;
typedef long long ll;
typedef long double ld;

vector <ll> vc;
ll t[N], cnt[N], ans;
int siz[N], i, n, u, v, w, kol;
vector <pair <int, ll> > g[N];
bool mk[N];

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

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

void addremains(int v, int pr, ll rem, ll remain)
{
    rem += cnt[v];

    vc.pb(remain);

    for (auto u : g[v])
    {
        int to = u.F;
        ll len = u.S;

        if (to == pr || mk[to]) continue;

        ll new_rem = max(0ll, rem - len );
        ll new_remain = remain + max(0ll, len - rem);

        addremains(to, v, new_rem, new_remain);
    }
}

void calcremains(int v, int pr, ll rem, ll remain)
{
    rem += cnt[v];

    upd(upper_bound(vc.begin(), vc.end(), remain) - vc.begin() - 1);

    for (auto u : g[v])
    {
        int to = u.F;
        ll len = u.S;

        if (to == pr || mk[to]) continue;

        ll new_rem = max(0ll, rem - len);
        ll new_remain = remain + max(0ll, len - rem);

        calcremains(to, v, new_rem, new_remain);
    }
}

void calcneeds(int v, int pr, ll rem, ll cost, ll need)
{
    siz[v] = 1;

    need += cost;
    need = max(0ll, need - cnt[v]);

    rem += cnt[v];
    rem -= cost;

    if (need == 0)
    {
        kol++;
        ans += sum(upper_bound(vc.begin(), vc.end(), rem) - vc.begin() - 1);
    }

    for (auto u : g[v])
    {
        int to = u.F;
        ll len = u.S;

        if (to == pr || mk[to]) continue;

        calcneeds(to, v, rem, len, need);

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

void dfs(int v, int pr)
{
    siz[v] = 1;
    for (auto u : g[v])
    {
        if (u.F == pr || mk[u.F]) continue;

        dfs(u.F, v);

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

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

    int cent = -1;

    for (auto u : g[v])
    {
        int to = u.F;

        if (mk[to]) continue;
        if (siz[to] > siz[v] / 2) {cent = to;break;}
    }

    if (cent == -1) return v;

    return fnd_cent(cent, v);
}

void calc(int v)
{
    kol = 0;
    mk[v] = 1;
    vc.clear();
    vc.pb(-1);

    for (auto u : g[v])
    {
        int to = u.F;
        ll len = u.S;

        if (mk[to]) continue;

        addremains(to, -1, 0, len);
    }

    sort(vc.begin(), vc.end());
    vc.resize(unique(vc.begin(), vc.end()) - vc.begin());

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

    for (auto u : g[v])
    {
        int to = u.F;
        ll len = u.S;

        if (mk[to]) continue;

        calcneeds(to, -1, cnt[v], len, 0);

        calcremains(to, -1, 0, len);
    }

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

    for (int i = sz(g[v]) - 1; i >= 0; i--)
    {
        int to = g[v][i].F;
        ll len = g[v][i].S;

        if (mk[to]) continue;

        calcneeds(to, -1, cnt[v], len, 0);

        calcremains(to, -1, 0, len);
    }

    ans += kol / 2;

    ans += sum(upper_bound(vc.begin(), vc.end(), cnt[v]) - vc.begin() - 1);

    for (auto u : g[v])
    {
        int to = u.F;

        if (mk[to]) continue;

        dfs(to, -1);

        int cent = fnd_cent(to, -1);

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

int main()
{
    srand(time(0));
    ios_base::sync_with_stdio(0);
    iostream::sync_with_stdio(0);
    ios::sync_with_stdio(0);
    cin.tie(NULL);
    cout.tie(NULL);

//    in("input.txt");
//    out("output.txt");

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

    for (i = 1; i < n; i++)
    {
        cin >> v >> u >> w;

        g[u].pb({v, w});
        g[v].pb({u, w});
    }

    dfs(1, -1);

    int cent = fnd_cent(1, -1);

    calc(cent);

    cout << ans;
}

/*
8 8
...).).*
*....)..
.)*(..).
(*)((...
.)).)(..
.)(.)..(
...).(.*
M.......
*/
//
//110000
//1100

• 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