题目
题意:
给定一棵树,再给一个k,要求计算树上任意两点的距离除以k(向上取整)的和。
分析:
因为是树上问题,所以自然就想到子节点与父节点的转移。因为k比较小,所以我们维护一个num[i][j],表示到i节点的子树中到i的距离模k为j的节点数。那么对于子节点到父节点的转移,其实只要加上模k为0的那些点,到父节点需要加1的代价和当前子节点到父节点的代价,其余的只要用来更新父节点的num数组即可。这样我们就可以算出根节点的贡献。
由于要算其他的点的贡献,所以我们需要换根。换根的话就是将父节点的num数组减去当前子节点对它的贡献,然后子节点加上父节点的num数组,然后再算上父节点num数组模k为0的点的贡献和父节点的贡献(原先的父节点现在变成了子节点)。记得回溯时还原!!!
#include <iostream>
#include <vector>
using namespace std;
typedef long long ll;
ll dp[200005],num[200005][5];
int n,k;
vector<int> g[200005];
void dfs1(int x,int fa)
{
if( g[x].size() == 1 && g[x][0] == fa )
{
return;
}
for (int i = 0; i < g[x].size(); i++)
{
int t = g[x][i];
if( t == fa ) continue;
dfs1(t,x);
for (int j = 0; j < k; j++)
{
num[x][j] += num[t][(j-1+k)%k];
}
num[x][1%k] ++;
dp[x] += dp[t] + 1 + num[t][0];
}
}
void dfs2(int x,int fa)
{
for (int i = 0; i < g[x].size(); i++)
{
int t = g[x][i];
if( t == fa ) continue;
for (int j = 0; j < k; j++)
{
num[x][j] -= num[t][(j-1+k)%k];
}
num[x][1%k] --;
dp[t] += dp[x] - dp[t] - 1 - num[t][0];
dp[t] += 1 + num[x][0];
for (int j = 0; j < k; j++)
{
num[t][j] += num[x][(j-1+k)%k];
}
num[t][1%k] ++;
dfs2(t,x);
num[t][1%k] --;
for (int j = 0; j < k; j++)
{
num[t][j] -= num[x][(j-1+k)%k];
}
num[x][1%k] ++;
for (int j = 0; j < k; j++)
{
num[x][j] += num[t][(j-1+k)%k];
}
}
}
int main()
{
ios::sync_with_stdio(false);
cin.tie(0);
cin >> n >> k;
for (int i = 1; i < n; i++)
{
int x,y;
cin >> x >> y;
g[x].push_back(y);
g[y].push_back(x);
}
dfs1(1,0);
dfs2(1,0);
ll ans = 0;
for (int i = 1; i <= n; i++)
{
ans += dp[i];
}
cout << ans / 2 << '\n';
return 0;
}