Submission #381914

#	Time	Username	Problem	Language	Result	Execution time	Memory
381914		kartel	Janjetina (COCI21_janjetina)	C++14	110 / 110	445 ms	17920 KiB

Main

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

const int N = 1e5 + 500;

vector <pii> g[N], vec, vc;

int siz[N], to[N], cntall, n, k, root, f[N];
ll ans;

void upd(int v, int val) {for (; v < N; v = (v | (v + 1))) f[v] += val;}
int get(int v) {int res = 0; for (; v >= 0; v = (v & (v + 1)) - 1) res += f[v]; return res;}

void calc_sizes(int v, int pr) {
    siz[v] = 1;

    for (auto u : g[v]) {
        if (to[u.F] || u.F == pr) {
            continue;
        }

        calc_sizes(u.F, v);
        siz[v] += siz[u.F];
    }
}

int fnd_cent(int v, int pr) {
    int cent = -1;

    for (auto u : g[v]) {
        if (to[u.F] || u.F == pr) {
            continue;
        }

        if (siz[u.F] > siz[root] / 2 && (cent == -1 || siz[cent] < siz[u.F])) {
            cent = u.F;
        }
    }

    if (cent == -1) {
        return v;
    }
    else {
        return fnd_cent(cent, v);
    }
}

void calc(int v, int pr, int mx) {
    if (mx > k) {
        ans += cntall + 1;
    }

    for (auto u : g[v]) {
        if (to[u.F] || u.F == pr) {
            continue;
        }

        calc(u.F, v, max(mx, u.S));
    }
}

void add(int v, int pr, int mx, int l) {
    vc.pb({mx, l});
    vec.pb({mx, l});

    for (auto u : g[v]) {
        if (to[u.F] || u.F == pr) {
            continue;
        }

        add(u.F, v, max(mx, u.S), l + 1);
    }
}

void go(int v, int pr) {
    root = v;
    calc_sizes(v, pr);

    int cent = fnd_cent(v, pr);
    to[cent] = 1;

    vec.clear();

    cntall = 0;
    for (auto u : g[cent]) {
        if (to[u.F]) {
            continue;
        }

        vc.clear();

        add(u.F, cent, u.S, 1);

        sort(all(vc));

        for (auto x : vc) {
            int mx = x.F;
            int l = x.S;

            if (mx - l - k >= 0) {
                ans -= get(mx - l - k);
            }
             upd(l, 1);
        }

        for (auto x : vc) {
            upd(x.S, -1);
        }
    }

    sort(all(vec));

    for (int i = 0; i < sz(vec); i++) {
        int w1 = vec[i].F;
        int l1 = vec[i].S;

        if (w1 - l1 >= k) {
            ans++;
        }

        if (w1 - l1 - k >= 0) {
            ans += get(w1 - l1 - k);
        }

        upd(l1, 1);
    }
    for (int i = 0; i < sz(vec); i++) {
        upd(vec[i].S, -1);
    }

    for (auto u : g[cent]) {
        if (to[u.F]) {
            continue;
        }

        go(u.F, v);
    }
}

int main()
{
//    mt19937 rnd(chrono::steady_clock::now().time_since_epoch().count());;
	ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
//	in("toys.in");
//	out("toys.out");
//    in("input.txt");
//    out("output.txt");
//    cerr.precision(9); cerr << fixed;

//    clock_t tStart = clock();

    cin >> n >> k;

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

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

    go(1, -1);

    cout << ans * 2ll;
}


/*

7
4 6 7 2 3 1 5
*/

• Subtask 1: 15 / 15
• Subtask 2: 35 / 35
• Subtask 3: 60 / 60