Submission #747909

#	Time	Username	Problem	Language	Result	Execution time	Memory
747909		FulopMate	Janjetina (COCI21_janjetina)	C++17	110 / 110	657 ms	19700 KiB

Main

#include <bits/stdc++.h>

// #define MULTI_TEST_CASE
// #define TEST_TEXT

using namespace std;

#define ll long long
#define MAX(a, b) (a) = max((a), (b))
#define MIN(a, b) (a) = min((a), (b))
#define all(a) (a).begin(), (a).end()
#define sortedpair(a, b) {min((a), (b)), max((a), (b))}

const ll MOD = 1e9+7;

struct edge {
    int to, w;
};

int n, k;
vector<vector<edge>> v;
ll ans = 0;
vector<bool> voltcentroid;
vector<int> siz;
const int st_size = 200001;
vector<ll> st;

void add(int d, int diff = 1) {
    d += st_size;
    st[d] += diff;
    for (d /= 2; d >= 1; d /= 2) {
        st[d] = st[2*d] + st[2*d+1];
    }
}

ll query(int l, int r) {
    l += st_size; r += st_size;
    ll res = 0;
    while (l <= r) {
        if (l % 2 == 1) res += st[l++];
        if (r % 2 == 0) res += st[r--];
        l /= 2; r /= 2;
    }
    return res;
}

// void add(vector<ll> &st, int d, int dif = 1, int ind = 1, int l = 0, int r = st_size-1){
//     if(l == r){
//         st[ind]+=dif;
//         return;
//     }
//     int mid = (l+r)/2;
//     if(d <= mid) add(st, d, dif, ind*2, l, mid);
//     else add(st, d, dif, ind*2+1, mid+1, r);
//     st[ind] = st[ind*2] + st[ind*2+1];
// }

// ll query(vector<ll> &st, int ql, int qr, int ind = 1, int l = 0, int r = st_size-1){
//     if(qr < ql)return 0;
//     if(r < ql || l > qr)return 0;
//     if(ql <= l && r <= qr) return st[ind];
//     int mid = (l+r)/2;
//     return query(st, ql, qr, ind*2, l, mid) + query(st, ql, qr, ind*2+1, mid+1, r);
// }

int c_szamoz(int x, int p = -1){
    siz[x] = 1;
    for(edge&i : v[x]){
        if(i.to != p && !voltcentroid[i.to])
            siz[x] += c_szamoz(i.to, x);
    }
    return siz[x];
}

int c_keres(int x, int meret, int p = -1){
    for(edge&i : v[x]){
        if(i.to != p && !voltcentroid[i.to]){
            if(siz[i.to] > meret/2){
                return c_keres(i.to, meret, x);
            }
        }
    }
    return x;
}

struct c_ut {
    int w, x, d;
    bool operator<(const c_ut & other) const {
        return -w < -other.w;
    }
};

vector<int> inserts;

void reset_st(){
    while(inserts.size()){
        add(inserts.back(), -1);
        inserts.pop_back();
    }
}

void megold(int x = 0){
    int meret = c_szamoz(x);
    int centroid = c_keres(x, meret);

    // st_ut_max.assign(st_size*2, 0);

    add(0, 1); inserts.push_back(0);

    vector<bool> seen(n, 0);
    seen[centroid] = 1;
    priority_queue<c_ut> q;
    for(edge&i : v[centroid]){
        if(!voltcentroid[i.to]){
            q.push({i.w, i.to, 1});
        }
    }
    while(!q.empty()){
        c_ut c = q.top();
        q.pop();
        seen[c.x] = 1;
        ans += query(0, c.w-c.d-k);
        add(c.d, 1); inserts.push_back(c.d);
        for(edge&i : v[c.x]){
            if(!seen[i.to] && !voltcentroid[i.to]){
                q.push({max(c.w, i.w), i.to, c.d+1});
            }
        }
    }

    reset_st();
    seen.assign(n, 0);
    for(edge&i : v[centroid]){
        if(!voltcentroid[i.to]){
            seen[centroid] = 1;
            priority_queue<c_ut> q;
            q.push({i.w, i.to, 1});
            while(!q.empty()){
                c_ut c = q.top();
                q.pop();
                if(seen[c.x])continue;
                seen[c.x] = 1;
                ans -= query(0, c.w-c.d-k);
                add(c.d, 1); inserts.push_back(c.d);
                for(edge&i : v[c.x]){
                    if(!seen[i.to] && !voltcentroid[i.to]){
                        q.push({max(c.w, i.w), i.to, c.d+1});
                    }
                }
            }
            reset_st();
        }
    }


    // add(st_ut, 0);

    // for(edge&i : v[centroid]){
    //     if(!voltcentroid[i.to]){
    //         c_szamol(i.to, i.w, 1, centroid);
    //         c_frissit(i.to, i.w, 1, centroid);
    //     }
    // }


    voltcentroid[centroid] = true;
    for(edge&i : v[centroid]){
        if(!voltcentroid[i.to]){
            megold(i.to);
        }
    }
}

void solve(){
    cin>>n>>k;
    v.assign(n, {});
    voltcentroid.assign(n, 0);
    st.assign(2*st_size, 0);
    siz.assign(n, 0);
    for(int i = 0; i < n-1; i++){
        int a, b, c; cin>>a>>b>>c; a--; b--;
        v[a].push_back({b, c});
        v[b].push_back({a, c});
    }
    megold();
    cout<<ans*2<<endl;
}

int main(){
    ios_base::sync_with_stdio(0); cin.tie(0);
    int _t = 1;
#ifdef MULTI_TEST_CASE
    cin >> _t;
#endif
    for(int _i = 0; _i < _t; _i++){
        #ifdef TEST_TEXT
        cout<<"Case #"<<_i+1<<": ";
        #endif
        solve();
    }
    return 0;
}

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