Centroid Decomposition - YessineJallouli/Competitive-Programming GitHub Wiki

#include <bits/stdc++.h>
typedef long long ll;
using namespace std;
 
const int N = 2e5+7;
bool marked[N];
int subtree_size[N];
vector<int> tree[N];
int cnt[N]{1};
 
long long ans = 0;
int mxDepth = 0;
int k;
 
int getSubtreeSize(int node, int par = -1) {
    subtree_size[node] = 1;
    for (int i : tree[node]) {
        if (i != par && ! marked[i]) {
            subtree_size[node]+= getSubtreeSize(i, node);
        }
    }
    return subtree_size[node];
}
 
int get_centroid(int node, int sz, int par = -1) {
    for (int i : tree[node]) {
        if (i != par && ! marked[i]) {
            if (subtree_size[i]*2 > sz) {
                return get_centroid(i, sz, node);
            }
        }
    }
    return node;
}
 
void get_cnt(int node, bool filling, int par, int depth = 1) {
    if (depth > k)
        return;
    mxDepth = max(depth, mxDepth);
    if (filling)
        cnt[depth]++;
    else
        ans+= cnt[k-depth];
    for (int i : tree[node]) {
        if (i != par && ! marked[i]) {
            get_cnt(i, filling, node, depth+1);
        }
    }
}
 
void centroid(int node = 1) {
    int c = get_centroid(node, getSubtreeSize(node));
    mxDepth = 1;
    for (int i : tree[c]) {
        if (! marked[i]) {
            get_cnt(i, false, c);
            get_cnt(i, true, c);
        }
    }
    marked[c] = true;
    fill(cnt+1, cnt+mxDepth+1, 0);
    for (int i : tree[c]) {
        if (! marked[i]) {
            centroid(i);
        }
    }
}
 
 
int main() {
    int n;
    cin >> n >> k;
    for (int i = 0; i < n-1; i++) {
        int a,b; cin >> a >> b;
        tree[a].push_back(b);
        tree[b].push_back(a);
    }
    centroid();
    cout << ans;
}

#include <bits/stdc++.h>
typedef long long ll;
using namespace std;
 
const int N = 2e5+7;
bool marked[N];
int subtree_size[N];
vector<int> tree[N];
 
ll fenTree[N];
void upd(int x, ll val) {
    for (; x < N; x+=x&(-x)) {
        fenTree[x]+= val;
    }
}
 
ll get(int x) {
    ll res = 0;
    for (; x; x-=x&(-x)) {
        res+= fenTree[x];
    }
    return res;
}
 
ll getRange(int a, int b) {
    if (a == 0)
        return get(b);
    return get(b) - get(a-1);
}
 
long long ans = 0;
int mxDepth = 0;
int l,r;
 
int get_subtree_sizes(int node, int parent = -1) {
    subtree_size[node] = 1;
    for (int i : tree[node])
        if (!marked[i] && i != parent)
            subtree_size[node] += get_subtree_sizes(i, node);
    return subtree_size[node];
}
int get_centroid(int node, int sz, int par = -1) {
    for (int i : tree[node]) {
        if (i != par && ! marked[i]) {
            if (subtree_size[i]*2 > sz) {
                return get_centroid(i, sz, node);
            }
        }
    }
    return node;
}
 
void get_cnt(int node, bool filling, int par, int depth = 1) {
    if (depth > r)
        return;
    mxDepth = max(depth, mxDepth);
    if (filling) {
        upd(depth+1, 1);
    }
    else {
        ans+= getRange(max(1,l-depth+1), r-depth+1);
    }
    for (int i : tree[node]) {
        if (i != par && ! marked[i]) {
            get_cnt(i, filling, node, depth+1);
        }
    }
}
 
void centroid(int node = 1) {
    int c = get_centroid(node, get_subtree_sizes(node));
    mxDepth = 1;
    for (int i : tree[c]) {
        if (! marked[i]) {
            get_cnt(i, false, c);
            get_cnt(i, true, c);
        }
    }
    marked[c] = true;
    for (int i = 1; i <= mxDepth+1; i++) {
        ll val = getRange(i,i);
        upd(i,-val);
    }
    upd(1, 1);
    for (int i : tree[c]) {
        if (! marked[i]) {
            centroid(i);
        }
    }
}
 
int main() {
    ios_base::sync_with_stdio(false), cin.tie(nullptr);
    int n;
    cin >> n >> l >> r;
    upd(1,1);
    for (int i = 0; i < n-1; i++) {
        int a,b; cin >> a >> b;
        tree[a].push_back(b);
        tree[b].push_back(a);
    }
    centroid();
    cout << ans;
}