Submission #541703

#	Time	Username	Problem	Language	Result	Execution time	Memory
541703		AlperenT	Janjetina (COCI21_janjetina)	C++17	110 / 110	404 ms	17532 KiB

Main

#include <bits/stdc++.h>
 
using namespace std;
 
const int N = 1e5 + 5;
 
int n, k, a, b, w, treesz[N];
 
bool isblocked[N];
 
struct Edge{
    int b, w;
};
 
vector<Edge> tree[N];
 
struct Item{
    int w, d;
 
    bool operator < (const Item &sc) const{
        if(w == sc.w) return d < sc.d;
        else return w < sc.w;
    }
};
 
struct Fenwick{
    int tree[N];
 
    stack<pair<int, int>> hist;
 
    int getsum(int r){
        int sum = 0;
        for(int i = r; i > 0; i -= i & -i) sum += tree[i];
        return sum;
    }
 
    int query(int l, int r){
        if(l > r) return 0;
        return getsum(r) - getsum(l - 1);
    }
 
    void update(int pos, int val, bool addhist = true){
        if(addhist) hist.push({pos, val});
        for(int i = pos; i < N; i += i & -i) tree[i] += val;
    }
 
    void clear(){
        while(!hist.empty()){
            update(hist.top().first, -hist.top().second, false);
            hist.pop();
        }
    }
};
 
Fenwick fw;
 
int calcsz(int v, int p){
    treesz[v] = 1;
 
    for(auto e : tree[v]){
        if(e.b != p && !isblocked[e.b]) treesz[v] += calcsz(e.b, v);
    }
 
    return treesz[v];
}
 
int findcent(int v, int p, int sz){
    for(auto e : tree[v]){
        if(e.b != p && !isblocked[e.b] && treesz[e.b] >= sz / 2){
            return findcent(e.b, v, sz);
        }
    }
 
    return v;
}
 
void dfs(int v, int p, int d, int w, vector<Item> &vec){
    vec.push_back({w, d});
 
    for(auto e : tree[v]){
        if(e.b != p && !isblocked[e.b]){
            dfs(e.b, v, d + 1, max(w, e.w), vec);
        }
    }
}
 
long long centdecomp(int v){
    int sz = calcsz(v, 0);
 
    int cent = findcent(v, 0, sz);
 
    long long ans = 0;
 
    vector<Item> total, cur;
 
    for(auto e : tree[cent]){
        if(!isblocked[e.b]){
            cur.clear();
 
            dfs(e.b, cent, 1, e.w, cur);
 
            sort(cur.begin(), cur.end());
 
            for(auto itm : cur){
                if(itm.w - itm.d >= k) ans += 1;
 
                ans -= fw.query(0, itm.w - itm.d - k);
 
                fw.update(itm.d, 1);
 
                total.push_back(itm);
            }
 
            fw.clear();
        }
    }
 
    sort(total.begin(), total.end());
 
    for(auto itm : total){
        ans += fw.query(0, itm.w - itm.d - k);
 
        fw.update(itm.d, 1);
    }
 
    fw.clear();
 
    isblocked[cent] = true;
 
    for(auto e : tree[cent]){
        if(!isblocked[e.b]){
            ans += centdecomp(e.b);
        }
    }
 
    return ans;
}
 
int main(){
    ios_base::sync_with_stdio(false);cin.tie(NULL);
 
    cin >> n >> k;
 
    for(int i = 1; i <= n - 1; i++){
        cin >> a >> b >> w;
 
        tree[a].push_back({b, w});
        tree[b].push_back({a, w});
    }
 
    cout << centdecomp(1) * 2;
}

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