codeforces-161-D : Distance in tree (点分治)

题目：https://codeforces.com/contest/161/problem/D

题意：求树上距离为k的点对数。

思路：点分治，参考博客：https://www.luogu.org/blog/user9012/dian-fen-zhi-lve-xie

代码：

#include <bits/stdc++.h>
using namespace std;
typedef pair<int, int> P;
const int maxn = 1e5+5;

vector<P> E[maxn];
int n, k, siz[maxn], rt, Siz, tot, ans, Mx_son[maxn], d[maxn], vis[maxn];
void get_root(int u, int fa){
    Mx_son[u] = 0; siz[u] = 1;
    for(auto it : E[u]){
        int v = it.first, w = it.second;
        if(v == fa || vis[v]) continue;
        get_root(v, u); 
        Mx_son[u] = max(Mx_son[u], siz[v]);
        siz[u] += siz[v];
    }
    Mx_son[u] = max(Mx_son[u], Siz-siz[u]);
    if(Mx_son[u] < Mx_son[rt]) rt = u;
}
void get_dis(int u, int fa, int D){
    for(auto it : E[u]){
        int v = it.first, w = it.second;
        if(v == fa || vis[v]) continue;
        d[++tot] = D+w;
        get_dis(v, u, d[tot]);
    }
}
int get_ans(int u, int D){
    d[tot=1] = D; 
    get_dis(u, 0, D);
    sort(d+1, d+tot+1);
    int l = 1, ans = 0;
    while(l<tot && d[l]+d[tot]<k) l++;
    while(l<tot && d[l]<=k-d[l]) {
        ans += ((upper_bound(d+l+1, d+tot+1, k-d[l])-d) - (lower_bound(d+l+1, d+tot+1, k-d[l])-d));
        l ++;
    }
    return ans;
}
void dfs(int u){
    vis[u] = 1; ans += get_ans(u, 0);
    for(auto it : E[u]){
        int v = it.first, w = it.second;
        if(vis[v]) continue;
        ans -= get_ans(v, w);
        Siz = siz[v]; rt = 0;
        get_root(v, u); dfs(rt);
    }
}
int main(){
    #ifndef ONLINE_JUDGE
        freopen("./data.in", "r", stdin);
        freopen("./data.out", "w", stdout);
    #endif
    int u, v;
    scanf("%d%d", &n, &k);
    for(int i=1; i<n; i++){
        scanf("%d%d", &u, &v);
        E[u].push_back(P(v, 1));
        E[v].push_back(P(u, 1));
    }
    Mx_son[0] = Siz = n;
    get_root(1, 0); dfs(rt);
    cout << ans << endl;
}