Submission #977495

#	Time	Username	Problem	Language	Result	Execution time	Memory
977495		tnknguyen_	Janjetina (COCI21_janjetina)	C++14	15 / 110	1558 ms	58704 KiB

Main

#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){
        while(i < n){
            f[i] += val;
            i += (i & -i);
        }
    }

    ll get(int i){
        ll ans = 0;
        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], id[MX];
int n, k;
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;
}

int idx = 0;
void solve(int u, int p, int len, int _max, vector<pair<pii, int>> &vals, const int &cent){
    vals.push_back({{_max, len}, u});
    id[u] = (p == cent ? ++idx : id[p]);

    for(pii e : gr[u]){
        int v, c;
        tie(v, c) = e;
        if(v != p && !del[v]){
            solve(v, u, len + 1, max(_max, c), vals, cent);
        }
    }
}

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


    vector<pair<pii, int>> vals;
    idx = 0;
    for(pii e : gr[u]){
        int v, c;
        tie(v, c) = e;
        if(!del[v]){
            solve(v, u, 1, c, vals, u);
        }
    }

    vector<BIT> cnt(idx + 1);
    sort(vals.begin(), vals.end());

    BIT bit;
    int res = 0;
    for(pair<pii, int> p : vals){
        int w, l, u = p.second;
        tie(w, l) = p.first;
        int need = w - l - k;
        // cout<<u<<' '<<w<<' '<<l<<"\n";
        // cout<<"need: "<<need<<' '<<bit.get(need) - cnt[id[u]].get(need)<<"\n\n";
        res += (bit.get(need) - cnt[id[u]].get(need));
        res += (w - l >= k);

        bit.update(l, 1);
        cnt[id[u]].update(l, 1);
    }
    // cout<<"Current Centroid: "<<u<<"\n";
    // cout<<res<<"\n";
    ans += res;

    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;
}

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