Submission #381946

#	Time	Username	Problem	Language	Result	Execution time	Memory
381946		VEGAnn	Janjetina (COCI21_janjetina)	C++14	110 / 110	1010 ms	21540 KiB

Main

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#define sz(x) ((int)x.size())
#define PB push_back
#define i2 array<int,2>
#define i3 array<int,3>
#define all(x) x.begin(),x.end()
using namespace std;
using namespace __gnu_pbds;
template <class T>
using ordered_set = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long long ll;
const int MX = int(1e7) + 10;
const int oo = 2e9;
const int N = 100100;
vector<i2> g[N], vc;
ordered_set<i2> st;
int n, k, siz[N], tt = 0;
bool mrk_cd[N];
ll ans;

void build_sz(int v, int p){
    siz[v] = 1;

    for (i2 U : g[v]){
        int u = U[0];

        if (u == p || mrk_cd[u]) continue;

        build_sz(u, v);

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

int search(int v, int p, int SZ){
    for (i2 U : g[v]){
        int u = U[0];

        if (u == p || mrk_cd[u]) continue;

        if (siz[u] > SZ / 2)
            return search(u, v, SZ);
    }

    return v;
}


void add(int v, int p, int mx, int len){
    vc.PB({mx, len});

    for (i2 u : g[v]){
        if (u[0] == p || mrk_cd[u[0]]) continue;

        add(u[0], v, max(mx, u[1]), len + 1);
    }
}

void build_cd(int v, int p){
    build_sz(v, -1);

    v = search(v, -1, siz[v]);

    mrk_cd[v] = 1;

    vc.clear();

    vc.PB({0, 0});

    for (i2 U : g[v]){
        int u = U[0];

        if (u == p || mrk_cd[u]) continue;

        add(u, v, U[1], 1);
    }

    sort(all(vc));
    st.clear();

    for (i2 cr : vc){
        int bd = cr[0] - k - cr[1];

        ans += st.order_of_key({bd, oo});

        st.insert({cr[1], ++tt});
    }

    for (i2 U : g[v]){
        int u = U[0];

        if (u == p || mrk_cd[u]) continue;

        vc.clear();

        add(u, v, U[1], 1);

        sort(all(vc));

        st.clear();

        for (i2 cr : vc){
            int bd = cr[0] - k - cr[1];

            ans -= st.order_of_key({bd, oo});

            st.insert({cr[1], ++tt});
        }
    }

    for (i2 U : g[v]){
        int u = U[0];

        if (u == p || mrk_cd[u]) continue;

        build_cd(u, v);
    }
}

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

#ifdef _LOCAL
    freopen("in.txt","r",stdin);
#endif // _LOCAL

    cin >> n >> k;

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

        x--; y--;

        g[x].PB({y, w});
        g[y].PB({x, w});
    }

    build_cd(0, -1);

    cout << ans * 2;

    return 0;
}

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