CF771C Bear and Tree Jumps

CF771C Bear and Tree Jumps

直接考虑换根。
但是有个K，不好搞，那就多存点，记录 \(f_{n,j}\) 为与 \(n\) 距离为 \(j\) 的点的答案（\(f_{n,0}\) 即为自己的答案），那么可以知道，对于距离不为 \(K\) 的倍数的点，可以直接距离加一，对于距离为 \(K\) 的倍数的点？对于 \(f_{n,0}\) 可以直接从 \(f_{v\in son_n,K-1}\) 转移，加上 \(siz_{n}-1\) 就好。
所以我们可以列出DP方程

\[\begin{matrix} f_{n,0} = \sum_{v\in son_n}f_{v,K-1} + siz_v \\ f_{n,i} = \sum_{v\in son_n}f_{v,i-1} * [i\ne 0] \end{matrix}\]

#include <iostream>
#include <cstring>
#include <cstdio>

using namespace std;

typedef long long ll;
const ll MAXN = 1e6+10;

struct edge {
    ll nt, to;
} E[MAXN];

ll N, head[MAXN], cnt = -1, ans = 0, K, f[MAXN][7], siz[MAXN];

void add(ll, ll);
void dfs(ll, ll);
void dfs2(ll, ll);

int main() {
    memset(head, -1, sizeof(head));
    scanf("%lld%lld", &N, &K);
    for (ll i = 1, x, y; i < N; i++) {
        scanf("%lld%lld", &x, &y);
        add(x, y);
        add(y, x);
    }
    dfs(1, 0);
    dfs2(1, 0);
    printf("%lld\n", ans / 2);
    return 0;
}

void dfs2(ll n, ll ff) {
    ll t[7] = {0};
    if (ff) {
        for (ll i = 1; i < K; i++) {
            t[i] = f[ff][i] - f[n][i-1];
        }
        t[0] = f[ff][0] - f[n][K-1] - siz[n];
        for (ll i = 1; i < K; i++) {
            f[n][i] += t[i-1];
        }
        f[n][0] += t[K-1] + (N - siz[n]);
    }
    ans += f[n][0];
    for (ll i = head[n]; ~i; i = E[i].nt) {
        ll v = E[i].to;
        if (v == ff) continue;
        else {
            dfs2(v, n);
        }
    }
}

void dfs(ll n, ll ff) {
    siz[n] = 1;
    for (ll i = head[n]; ~i; i = E[i].nt) {
        ll v = E[i].to;
        if (v == ff) continue;
        else {
            dfs(v, n);
            siz[n] += siz[v];
            f[n][0] += f[v][K-1] + siz[v];
            for (ll i = 1; i < K; i++) {
                f[n][i] += f[v][i-1];
            }
        }
    }
}

void add(ll x, ll y) {
    cnt++;
    E[cnt].to = y;
    E[cnt].nt = head[x];
    head[x] = cnt;
}

/*
6 2
1 2 
1 3 
2 4 
2 5
4 6
ans :20

6 1
1 2 
1 3 
2 4 
2 5
4 6
ans :31

7 2
1 2
2 3
3 4
4 5
5 6
6 7
ans :34
*/