Submission #747892

#	Time	Username	Problem	Language	Result	Execution time	Memory
747892		FulopMate	Janjetina (COCI21_janjetina)	C++17	15 / 110	1559 ms	5716 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 = 262144;
vector<ll> st_ut, st_ut_max;

void add(vector<ll> &st, int d, int ind = 1, int l = 0, int r = st_size-1){
    if(l == r){
        st[ind]++;
        return;
    }
    int mid = (l+r)/2;
    if(d <= mid) add(st, d, ind*2, l, mid);
    else add(st, d, 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){
    if(voltcentroid[x])return 0;
    int cur = 1;
    for(edge&i : v[x]){
        if(i.to == p)continue;
        cur += c_szamoz(i.to, x);
    }
    return siz[x] = cur;
}

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

// void c_szamol(int x, int mx_w, int d, int p = -1){
//     ans += query(st_ut, 0, mx_w-d-k);
//     ans += query(st_ut_max, d+k, 2001);
//     for(edge&i : v[x]){
//         if(i.to == p)continue;
//         if(!voltcentroid[i.to]){
//             c_szamol(i.to, max(mx_w, i.w), d+1, x);
//         }
//     }
// }

// void c_frissit(int x, int mx_w, int d, int p = -1){
//     add(st_ut, d);
//     if(mx_w >= d){
//         add(st_ut_max, min(st_size-1, mx_w-d));
//     }
//     for(edge&i : v[x]){
//         if(i.to == p)continue;
//         if(!voltcentroid[i.to]){
//             c_frissit(i.to, max(mx_w, i.w), d+1, x);
//         }
//     }
// }

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

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

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

    add(st_ut, 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();
        if(seen[c.x])continue;
        seen[c.x] = 1;
        ans += query(st_ut, 0, c.w-c.d-k);
        add(st_ut, 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});
            }
        }
    }

    for(edge&i : v[centroid]){
        if(!voltcentroid[i.to]){
            st_ut.assign(st_size*2, 0);
            // add(st_ut, 0);
            seen.assign(n, 0);
            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(st_ut, 0, c.w-c.d-k);
                add(st_ut, 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});
                    }
                }
            }

        }
    }


    // 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);
    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: 0 / 35
• Subtask 3: 0 / 60