codeforces 791D 树形dp 所有点对

题意：20w个点的树，边长都为1，每个点可以往距离不超过k<=5的点进行一次跳跃。求所有的f(A,B)的和，f(A,B)是A到B所需的最小跳跃次数，标号A<B

题解：树形dp，A<B的条件就是指每个点对只扫一次，每次记录到每个点长度%k为0~k-1的个数以及有多少整段的跳跃。然后进行转移即可。对于每个点对只扫一遍的原理：每个点只需与父节点已扫的子树进行结合，然后自己再加入父节点的当前状态，对其进行更新。

#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
ll dp[250000][5];
vector<int>g[250000];
int k;
ll dists[220000],ans,sons[220000];
void dfs(int x,int fa)
{
	sons[x]=1;
	dp[x][0]=1;
	int u;
	for(int t=0;t<g[x].size();t++)
    {
        u=g[x][t];
        if(u==fa) continue;
        dfs(u,x);
        ans+=sons[x]*(dists[u])+dists[x]*sons[u];
        for(int i=0;i<k;i++)
            for(int j=0;j<k;j++)
            {
				ans+=dp[x][i]*dp[u][j];
				if(j+i>=k)ans+=dp[x][i]*dp[u][j];
			}
        sons[x]+=sons[u];
        dists[x]+=dists[u];
        dists[x]+=dp[u][k-1];
            for(int i=0;i<k;i++) dp[x][(i+1)%k]+=dp[u][i];
    }
}
int main(){
	int n;
	cin>>n>>k;ans=0;
	for(int i=1;i<n;i++)
    {
		int a,b;
		cin>>a>>b;
		g[a].push_back(b);
		g[b].push_back(a);
	}
	dfs(1,0);
	cout<<ans<<endl;
}