[换根dp] CF771C

_Cheems · 2024-01-27 20:53:50 · 题解

一个稍有不同的解法？

题意

给出一棵有 n 个点的树，定义 dis(i,j) 为树上两点距离，求 \sum\limits_{1\le i < j \le n} \lceil \frac {dis(i,j)}{k} \rceil。

思路

容易想到换根 \rm dp 求出每个点到其它点的距离之和 S_i，然后除以 k，但是由于有向上取整所以不好处理。

如何解决呢？发现不好对每个点单独讨论，也就是说最终式子最好含有 S_i 项。

我们知道 \lceil \frac ak \rceil = \frac {a+k-a \bmod k}{k}，只要把 a 除不尽的部分补上即可（当然前提是 k\nmid a ）。

根据上式，有：\sum \lceil \frac {a_i}{k} \rceil = \sum \frac {a_i+k-a_i \bmod k}{k}=\frac {\sum a_i + \sum (k-a_i\bmod k)} {k}。

诶，这样就出现 S_i 项了啊。把刚刚推出的式子带到这题里，节点 i 对答案的贡献就是：\frac {S_i+\sum k-dis(i,j) \bmod k}k 。

考虑怎么求 \sum k-dis(i,j) \bmod k 这一项即可。注意到 k 很小，容易想到枚举余数 c，统计 dep(i,j) \bmod k = c 的数量即可。

具体的，记 g_{i,c}/f_{i,c} 为子树内/子树外有多少个点 j 满足 dep(i,j)\bmod k=c，直接维护即可。

复杂度 O(nk)，可以通过。最后注意：上述过程统计的是有序数对，转成无序数对要除以 2。

代码

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

#define int long long
const int N = 2e5 + 5, K = 5;
int n, k, g[N], f[N], son[N], u, v;
int G[N][K], F[N][K], Gs[N][K], ans;
vector<int> to[N];

inline void dfs1(int u, int fa){
    son[u] = 1;
    for(auto v : to[u])
        if(v ^ fa){
            dfs1(v, u);
            son[u] += son[v];
            g[u] += g[v] + son[v];

            for(int c = 0; c < k; ++c)
                G[u][(c + 1) % k] += G[v][c], Gs[u][c] += G[v][c];
            ++G[u][1 % k], ++Gs[u][0];
        }
}
inline void dfs2(int u, int fa){
    if(fa){
        f[u] = f[fa] + n - son[fa] + 1;
        f[u] += g[fa] - g[u] - son[u] + son[fa] - son[u] - 1;

        for(int c = 0; c < k; ++c){
            F[u][(c + 1) % k] += F[fa][c];
            F[u][(c + 2) % k] += Gs[fa][c] - G[u][c] - (c == 0);
        }
        ++F[u][1 % k];
    }
    for(auto v : to[u])
        if(v ^ fa) dfs2(v, u);
}
signed main(){
    cin >> n >> k;
    for(int i = 1; i < n; ++i)
        scanf("%lld%lld", &u, &v), to[u].push_back(v), to[v].push_back(u);
    dfs1(1, 0), dfs2(1, 0);
    for(int i = 1; i <= n; ++i){
        int res = g[i] + f[i];
        for(int c = 1; c < k; ++c)
            res += (k - c) * (G[i][c] + F[i][c]);
        ans += res / k;
    }
    cout << ans / 2;
    return 0;
}