[codeforces161D]Distance in Tree(点分治/树形dp)

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

解题关键：练习一下点分治不用容斥 而直接做的做法。注意先查询，后更新。

不过这个方法有个缺陷，每次以一个新节点为根，必须memset mp数组，或许使用map会好些，更新序号一类用ca这种形式更好些。

试了一下，map更慢，应该是带log的原因。

点分治解法：

 #pragma comment(linker,"/STACK:102400000,102400000")
 #include<cstdio>
 #include<cstring>
 #include<algorithm>
 #include<cstdlib>
 #include<iostream>
 #include<cmath>
 #include<map>
 #define maxn 100040
 #define maxm 1000500
 using namespace std;
 typedef long long ll;
 const ll mod=;
 const ll inf=1ll<<;
 ll n,k,ans,size,s[maxn],f[maxn],path[maxn],cr;
 ll head[maxn],cnt,root;
 bool vis[maxn];
 struct edge{
     ll to,nxt;
 }e[maxn<<];
 map<int,int>mp;
 void add_edge(ll u,ll v){
     e[cnt].to=v;
     e[cnt].nxt=head[u];
     head[u]=cnt++;
 }

 inline ll read(){
     char k=;char ls;ls=getchar();for(;ls<''||ls>'';k=ls,ls=getchar());
     ll x=;for(;ls>=''&&ls<='';ls=getchar())x=(x<<)+(x<<)+ls-'';
     if(k=='-')x=-x;return x;
 }

 void get_root(ll u,ll fa){//get_root会用到size
     s[u]=;f[u]=;//f是dp数组
     for(ll i=head[u];i!=-;i=e[i].nxt){
         ll v=e[i].to;
         if(v==fa||vis[v]) continue;
         get_root(v,u);
         s[u]+=s[v];
         f[u]=max(f[u],s[v]);
     }
     f[u]=max(f[u],size-s[u]);
     root=f[root]>f[u]?u:root;
 }

 void get_path_size(ll u,ll fa,ll dis){
     if(dis+<=k){
         path[cr]=dis+;
         cr++;
     }
     s[u]=;
     for(ll i=head[u];i!=-;i=e[i].nxt){
         ll v=e[i].to;
         if(v==fa||vis[v]) continue;
         get_path_size(v,u,dis+);
         s[u]+=s[v];
     }
 }

 void work(ll u,ll fa){
     vis[u]=true;
     mp.clear();
     mp[]=;
     for(ll i=head[u];i!=-;i=e[i].nxt){
         ll v=e[i].to;
         if(v==fa||vis[v]) continue;
         cr=;
         get_path_size(v,u,);
         for(ll j=;j<cr;j++){
             ans+=mp[k-path[j]];
         }
         for(int j=;j<cr;j++){
             mp[path[j]]++;
         }
     }
     for(ll i=head[u];i!=-;i=e[i].nxt){
         ll v=e[i].to;
         if(vis[v]||v==fa) continue;
         size=s[v],root=;
         get_root(v,u);
         work(root,u);
     }
 }

 void init(){
     memset(vis,,sizeof vis);
     memset(head,-,sizeof head);
     ans=cnt=;
 }

 int main(){
     ll a,b;
     f[]=inf;
     while(scanf("%I64d%I64d",&n,&k)!=EOF){
         init();
         for(int i=;i<n-;i++){
             a=read(),b=read();
             add_edge(a,b);
             add_edge(b,a);
         }
         size=n,root=;
         get_root(,-);
         work(root,-);
         printf("%d\n",ans);
     }
     return ;
 }

树形dp解法：

复杂度：$O(nk)$

 #include<bits/stdc++.h>
 using namespace std;
 typedef long long ll;
 #define maxn 100006
 int head[maxn],cnt,dp[maxn][];
 struct edge{
     int to,w,nxt;
 }e[maxn<<];
 ll ans;
 int n,k,a,b;
 void add_edge(int u,int v){
     e[cnt].to=v;
     e[cnt].nxt=head[u];
     head[u]=cnt++;
 }

 void dfs(int u,int fa){
     dp[u][]=;
     for(int i=head[u];i!=-;i=e[i].nxt){
         int v=e[i].to;
         if(v==fa) continue;
         dfs(v,u);
         for(int j=;j<=k;j++){
             dp[u][j]+=dp[v][j-];
         }
     }
     ans+=dp[u][k];
     int tmp=;
     for(int i=head[u];i!=-;i=e[i].nxt){
         int v=e[i].to;
         if(v==fa) continue;
         for(int j=;j<k;j++){
             tmp+=1ll*dp[v][j-]*(dp[u][k-j]-dp[v][k-j-]);
         }
     }
     ans+=tmp/;
 }

 void init(){
     memset(head,-,sizeof head);
     cnt=;
     ans=;
 }

 int main(){
     init();
     cin>>n>>k;
     for(int i=;i<n-;i++){
         cin>>a>>b;
         add_edge(a,b);
         add_edge(b,a);
     }
     dfs(,-);
     cout<<ans<<"\n";
     return ;
 }