Submission #977501

#TimeUsernameProblemLanguageResultExecution timeMemory
977501tnknguyen_Janjetina (COCI21_janjetina)C++14
110 / 110
381 ms16740 KiB
/**
 * @file main.cpp
 * @author tnknguyen_ (nguyenngockhoitran@gmail.com)
 * @date 08-05-2024
 * 
 * @copyright Copyright (c) 2024
 * https://oj.vnoi.info/problem/coci2021_r4_janjetina
 */
#include <bits/stdc++.h>
using namespace std;
#define endl '\n' 
#define ll long long 
#define len(s) (int)s.size() 
#define pii pair<int, int>
//#define int long long 

struct BIT{
    int n;
    vector<ll> f;

    BIT(int N = 1e5): n(N+5), f(N+5, 0) {};

    void update(int i, int val){
        ++i;
        while(i < n){
            f[i] += val;
            i += (i & -i);
        }
    }

    ll get(int i){
        ll ans = 0;
        ++i;
        while(i > 0){
            ans += f[i];
            i -= (i & -i);
        }
        return ans;
    }

    ll get(int l, int r){
        return get(r) - get(l-1);
    }
};

const int MX = 1e5 + 5;
vector<pii> gr[MX];
int del[MX], sz[MX];
int n, k;
BIT bit;
ll ans = 0;

int dfs(int u, int p){
    sz[u] = 1;
    for(pii e : gr[u]){
        int v = e.first;
        if(v != p && !del[v]){
            sz[u] += dfs(v, u);
        }
    }
    return sz[u];
}

int centroid(int u, int p, const int &n){
    for(pii e : gr[u]){
        int v = e.first;
        if(v != p && !del[v] && sz[v]*2 > n){
            return centroid(v, u, n);
        }
    }
    return u;
}

vector<pii> vals;
void calc(int u, int p, int len, int _max){ // Calculates vals[].
    vals.push_back({_max, len});
    for(pii e : gr[u]){
        int v, c;
        tie(v, c) = e;
        if(v != p && !del[v]){
            calc(v, u, len + 1, max(_max, c));
        }
    }
}

void solve(const int &op){ //update answer
    sort(vals.begin(), vals.end());

    for(pii p : vals){
        int w, l;
        tie(w, l) = p;
        int need = w - l - k;
        ans += (bit.get(need) * op);
        bit.update(l, 1);
    }

    for(pii p : vals){
        bit.update(p.second, -1);
    }
}

void decompose(int u){ // Centroid Decomposition
    dfs(u, 0);
    u = centroid(u, 0, sz[u]);
    del[u] = 1;

    vals.clear();
    calc(u, 0, 0, 0);
    solve(1);

    for(pii e : gr[u]){
        int v, c;
        tie(v, c) = e;
        if(!del[v]){
            vals.clear();
            calc(v, 0, 1, c);
            solve(-1);
        }
    }

    for(pii e : gr[u]){
        int v = e.first;
        if(!del[v]){
            decompose(v);
        }
    }
}

int32_t main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

//    freopen("main.inp","r",stdin);
//    freopen("main.out","w",stdout);

    cin>>n>>k;

    for(int i=1; i<n; ++i){
        int u, v, w;
        cin>>u>>v>>w;
        gr[u].push_back({v, w});
        gr[v].push_back({u, w});
    }

    decompose(1);
    cout<<ans*2;

    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...