[CF771C] Bear and Tree Jumps

[CF771C] Bear and Tree Jumps - 树形dp

Description

有一棵 n 个结点且有 n-1 条边的无根树，一只熊可以从当前节点可以跳到任何与当前节点距离不超过 k ((k le 5)) 的节点，设 f(u,v) 表示从 u 到 v 的最少跳跃次数。问对树上所有点对，f(u,v) 和为多少。

Solution

所求就是树上所有点对距离除以 k 上取整的总和

考虑将上取整多算的那部分加上去

dp，设 (f[p][i]) 表示 p 为根的子树中，到 p 距离取模后为 i 的点的数目，那么我们就可以在 LCA 处归并出路径，最后得到 (h[i]) 表示路径长度取模后为 i 的路径的数目

对于 i=0 的，我们不需要修正；其余的每条路径需要在总和中加上 k-i

#include <bits/stdc++.h>
using namespace std;

#define int long long 
const int N = 1000005;

vector <int >g[N];

int f[N][6], h[N], n,k,ans,siz[N];

void dfs(int p,int from)
{
    siz[p]=1;
    f[p][0]=1;
    for(int q:g[p]) if(q!=from) 
    {
        dfs(q,p);
        siz[p]+=siz[q];
        for(int i=0;i<k;i++) 
        {
            for(int j=0;j<k;j++)
            {
                h[(i+j+1)%k]+=f[p][i]*f[q][j];
            }
            f[p][i]+=f[q][(i-1+k)%k];
        }
    }
}


signed main()
{
    ios::sync_with_stdio(false);

    cin>>n>>k;
    for(int i=1;i<n;i++) 
    {
        int u,v;
        cin>>u>>v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1,0);
    for(int i=1;i<=n;i++) ans+=siz[i]*(n-siz[i]);
    for(int i=1;i<k;i++) ans+=h[i]*(k-i);
    cout<<ans/k<<endl;
}