点分治在树中的应用
看了论文才会写的https://wenku.baidu.com/view/1bc2e4ea172ded630b1cb602.html
题意:统计树中节点距离dis(i,j)<=k的对数
解法:由于暴力dfs为O(n^2),dp为O(n*k),故采用效率较高的分治算法。
记dis[I]为节点i到根的距离 所以dis[i]+dis[j]<=k 的ij对数就是答案,排序之后O(n)即可算出
由于递归计算子树时出现重复,所以要减去所有在同一子树中的ij对数,同样是递归求解。
//
// main.cpp
// poj1741
//
// Created by Charles on 2017/10/3.
// Copyright © 2017年 Charles. All rights reserved.
//
#include <iostream>
#include <vector>
#include <stdio.h>
#include <cstring>
#include <algorithm>
using namespace std;
const int maxn=1e4+5;
struct node{
int v,l;
node(int a,int b){v=a;l=b;};
};
vector<node> g[maxn];
vector<int> deep;
int n,k,sons[maxn],maxson[maxn],vis[maxn],dis[maxn],size,root,ans;
void getroot(int x,int fa){
int u;
sons[x]=1; maxson[x]=0;
for(int i=0;i<g[x].size();++i)
if((u=g[x][i].v)!=fa && !vis[u]){
getroot(u, x);
sons[x]+=sons[u];
maxson[x]=max(maxson[x],sons[u]);
}
maxson[x]=max(maxson[x],size-sons[x]);
if(maxson[x]<maxson[root]) root=x;
}
void getdeep(int x,int fa){
int u;
deep.push_back(dis[x]);
sons[x]=1;
for(int i=0;i<g[x].size();++i)
if((u=g[x][i].v)!=fa && !vis[u]){
dis[u]=dis[x]+g[x][i].l;
getdeep(u, x);
sons[x]+=sons[u];
}
}
int cal(int x,int init){
deep.clear();dis[x]=init;
getdeep(x, 0);
sort(deep.begin(), deep.end());
int ans=0;
for (int l=0, r=(int)deep.size()-1; l<r; )
if (deep[l] + deep[r] <= k) ans += r-l++;
else r--;
return ans;
}
void solve(int x){
int u;
ans+=cal(x,0);
vis[x]=1;
for(int i=0;i<g[x].size();++i)
if(!vis[u=g[x][i].v]){
ans-=cal(u, g[x][i].l);
maxson[0]=size=sons[u];
getroot(u, root=0);
solve(root);
}
}
int main(int argc, const char * argv[]) {
int a,b,c;
while(~scanf("%d%d",&n,&k) && n+k){
for(int i=0;i<=n;++i) g[i].clear();
memset(vis, 0, sizeof(vis));
for(int i=0;i<n-1;++i) {
scanf("%d%d%d",&a,&b,&c);
g[a].push_back(node(b,c));
g[b].push_back(node(a,c));
}
maxson[0]=size=n;
root=0;
getroot(1, root);
ans=0;
solve(root);
cout<<ans<<endl;
}
return 0;
}