[换根dp] CF771C

posted on 2024-01-27 12:53:50 | under 题集 | source

一个稍有不同的解法？

题意

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

\(n\le200000,1\le k\le5\)

思路

容易想到换根 \(\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;
}