Submission #787101

#	Time	Username	Problem	Language	Result	Execution time	Memory
787101		MISM06	Janjetina (COCI21_janjetina)	C++14	0 / 110	2 ms	4948 KiB

Main

//0 1 1 0 1
//0 1 0 0 1
//1 0 0 1 1
//0 1 1 0 1
#include <bits/stdc++.h>
#pragma GCC optimize("Ofast,unroll-loops")
// #pragma GCC target("avx2")

using namespace std;

#define F 			first
#define S 			second
#define pb 			push_back
#define sze			size()
#define	all(x)		x.begin() , x.end()
#define wall__		cout << "--------------------------------------\n";
#define kids		int mid = (tl + tr) >> 1, cl = v << 1, cr = v << 1 | 1
#define file_io		freopen("input.cpp", "r", stdin); freopen("output.cpp", "w", stdout);

typedef long long ll;
typedef long double dl;
typedef pair < int , int > pii;
typedef pair < int , ll > pil;
typedef pair < ll , int > pli;
typedef pair < ll , ll > pll;
typedef pair < int , pii > piii;
typedef pair < ll, pll > plll;


const ll N = 2e5 + 10;
const ll mod = 1e9 + 7;
const ll inf = 2e16;
const ll INF = 1e9 + 10;
const ll lg = 32;

int n, k;
vector < pii > g[N], ord;

int sub[N], flag[N], bit[N];
pii h[N];
ll ans = 0;

void pre_calc (int v, int p) {
    sub[v] = 1;
    for (auto e : g[v]) {
        int u = e.F;
        if (flag[u] || u == p) continue;
        pre_calc(u, v);
        sub[v] += sub[u];
    }
}
int find_cen (int v, int p, int s) {
    for (auto e : g[v]) {
        int u = e.F;
        if (u == p || flag[u]) continue;
        if (sub[u] > s / 2) return find_cen(u, v, s);
    }
    return v;
}
void dfs (int v, int p) {
    // cout << v << " , " << p << '\n';
    sub[v] = 1;
    if (p != 0) ord.pb({h[v].S, h[v].F});
    for (auto e : g[v]) {
        int u = e.F;
        if (flag[u] || u == p) continue;
        h[u] = {h[v].F + 1, max(h[v].S, e.S)};
        dfs(u, v);
        sub[u] += sub[v];
    }
}

inline int read (int idx) {
    int res = 0;
    for (; idx > 0; idx -= (idx & -idx)) res += bit[idx];
    return res;
}
inline void upd (int idx, int val, int s) {
    for (; idx <= s; idx += (idx & -idx)) bit[idx] += val;
}
ll calc (int root, int t) {
    // cout << root << " : \n";
    ord.clear();
    int pp = 0;
    if (t) pp = -1;
    dfs (root, pp);
    for (int i = 0; i <= sub[root]; i++) bit[i] = 0;
    sort(all(ord));
    // cout << "& \n";
    ll res = 0;
    for (auto e : ord) {
        // cout << e.F << " - " << e.S << '\n';
        int w = e.F, l = e.S;
        int l2 = w - k - l;
        if (l2 >= 0) {
            res += read(l2);
            if (t == 0) ++res;
        }
        if (l > 0) upd(l, 1, sub[root]);
    }
    // cout << "-- " << res << '\n';
    // wall__
    return res;
}
void centree (int v) {
    // cout << v << '\n';
    pre_calc(v, 0);
    // cout << "*\n";
    int cen = find_cen(v, 0, sub[v]);
    // cout << "cen = " << cen << '\n';
    h[cen] = {0, 0};
    ans += calc(cen, 0);
    // cout << "* " << ans << '\n';
    // cout << "* =\n";
    flag[cen] = 1;
    for (auto e : g[cen]) {
        int u = e.F;
        if (flag[u]) continue;
        h[u] = {1, e.S};
        ans -= calc(u, 1);
        // cout << "* " << ans << '\n';
        centree(u);
    }
}

void solve () {

	cin >> n >> k;
    for (int i = 1; i < n; i++) {
        int v, u, w; cin >> v >> u >> w;
        g[v].pb({u, w});
        g[u].pb({v, w}); 
    }
    centree(1);
    cout << ans * 2ll << '\n';
}


int main() {
	// ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
	int t = 1;
	// cin >> t;
	while (t--) {solve();}
    return 0;
}
/*
*/
//shrek is love;

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