codeforces-771C-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;
}