题目: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;
}