CF791D-Bear and Tree Jumps

要点在注释里，写题确实应该先想，如果想的时间够多了，写起来很快。

#include <iostream>
#include <cstdio>
#include <vector>
using namespace std;
int N,k;
long long ans = 0; // 答案会爆掉 long long
int sum[222222],num[222222][6],in[222222];
//sum子树节点个数 num[u][i] u子树中，距离根节点距离mod k == i 的节点个数
vector<int> E[222222];
void dfs(int u,int depth) //深度用来求两片区域相隔的距离
{
sum[u] = num[u][depth%k] = 1;
for (unsigned int i = 0; i<E[u].size();i++)
{
int v = E[u][i];
if (in[v]==1) continue;
in[v] = 1;
dfs(v,depth+1);
for (int j = 0; j<k ; j++)
for (int kk = 0; kk<k; kk++)
{
int t = (k- ((( ( j + kk- (2*depth) ) % k ) + k ) % k )) % k;
ans+=t*(long long)num[u][j]*(long long)num[v][kk]; // 将走的过程中 多余（浪费）的边计数
}
for (int j = 0; j<k; j++)
num[u][j]+=num[v][j];
sum[u]+=sum[v];
}
ans+=(long long)sum[u]*(long long)(N-sum[u]); // 统计每条边需要走的次数
}
int main()
{
cin >> N >> k;
for (int i = 1; i<N; i++)
{
int a,b;
scanf("%d%d",&a,&b);
E[a].push_back(b);
E[b].push_back(a);
}
in[1] = 1;
dfs(1,0);
cout << ans/k <<endl;
}