Meaning

• 一个长度为 $n$ 的序列，每个位置可以被染成 $[1,m]$ 中的任一颜色
• 求所有 $m^n$ 种染色方案的最长连续同色子段长度之和
• 答案对 $p$ 取模
• $1\le n\le300000$
• $2\le m\le10^8$
• $0.99\times10^9\le p\le1.01\times10^9$
• $p$ 为质数
• 时限 2s
• 空限 1G

Solution - Step 1

• 神仙 zzq 的题思路果然非常巧妙
• 我们知道，对于值域为非负整数的函数 $f(S)$ ，有
• $\sum_Sf(S)=\sum_{i\ge1}\sum_S[f(S)\ge i]$
• 回到问题。直觉告诉我们，这是与最值有关的计数问题
• 所以根据上面的式子进行转化可以简化问题
• $ans=\sum_Sf(S)=\sum_{i\ge1}\sum_S[f(S)\ge i]$
• 其中 $S$ 表示一种染色方案， $f(S)$ 表示染色方案 $S$ 的最长连续同色子段长度
• 容易得出 $\sum_S[f[S]\ge i]$ 表示最长连续同色子段长度 $\ge i$ 的方案数
• 「最大值 $\ge i$ 」还是不好处理，考虑补集转化一下
• $\sum_S[f(S)\ge i]=m^n-\sum_S[f(S)<i]$
• 于是
• $ans=nm^n\sum_{i=1}^n\sum_S[f(S)<i]=nm^n-\sum_{i=0}^{n-1}\sum_S[f(S)\le i]$
• 发现 $\sum_S[f[S]\le i]$ 实际上就是把长度为 $n$ 的序列分成若干段，每段长度都不超过 $i$ ，最后给每段染上 $[1,m]$ 内的颜色，相邻段不能染同色的方案数
• 先枚举段数 $j$
• 如果我们已经确定了划分方案，那么染色方案数显然是 $m(m-1)^{j-1}$
• 然后如果没有每段长度不超过 $i$ 的限制，那么划分方案数为 $\binom{n-1}{j-1}$
• 如果加上了限制，我们可以考虑容斥
• 即考虑 $j$ 段中的一部分段，让这些段的长度强行超过 $i$
• 也就是
• $\sum_{k=0}^j(-1)^k\binom jk\binom{n-ik-1}{j-1}$
• 其中 $(-1)^k$ 为容斥系数， $\binom jk$ 表示 $j$ 段中选 $k$ 段让这些段长度强行超过 $i$ ， $\binom{n-ik-1}{j-1}$ 表示 $j$ 段中的 $k$ 段已经默认长度超过 $i$ 的情况下，序列剩下的 $n-ik$ 个元素分给 $j$ 段的方案数
• 注意到
• $\binom {n-ik-1}{j-1}=\frac{(n-ik-1)!}{(j-1)!(n-ik-j)!}$
• $=\frac j{n-ik}\times\frac{(n-ik)!}{j!(n-ik-j)!}=\frac j{n-ik}\binom{n-ik}j$
• 为什么我们要这样转化呢？ Step 2 会讲
• 于是
• $ans=nm^n-\sum_{i=0}^{n-1}\sum_{j=1}^nm(m-1)^{j-1}\sum_{k=0}^j\frac j{n-ik}(-1)^k\binom jk\binom{n-ik}j$
• $=nm^n-m\sum_{i=0}^{n-1}\sum_{k=0}^n\frac {(-1)^k}{n-ik}\sum_{j=k}^nj(m-1)^{j-1}\binom jk\binom{n-ik}j$
• 注意到 $\binom jk\binom{n-ik}j$ 会对答案贡献的条件为 $k\le j\le n-ik$
• 也就是 $k(i+1)\le n$
• 所以枚举 $i$ 之后只需将 $k$ 枚举到 $\lfloor\frac n{i+1}\rfloor$ 即可
• 这样如果不考虑 $\sum_{j=k}^n$ 及后面的东西，那么复杂度就是 $O(n\log n)$
• 我们的重点来了： $\sum_{j=k}^n$ 及后面的东西，我们能够 $O(1)$ 求吗？

Solution - Step 2

• 上面的问题，也就是求
• $\sum_{j=k}^nj(m-1)^{j-1}\binom jk\binom{n-ik}j$
• 直接化式子行不通，我们尝试考虑这个式子的组合意义
• 发现这个式子相当于在 $n-ik$ 个球中选一部分（设选出了 $j$ 个， $j\ge k$ ）（ $\binom {n-ik}{j}$ ），然后在这一部分球中再选出 $k$ 个（ $\binom jk$ ），第一次选出的 $j$ 个球中选出 $j-1$ 个球染成 $[1,m-1]$ 内的任一颜色（ $j(m-1)^{j-1}$ ），剩下的 $n-ik-j+1$ 个球全部染成颜色 $m$ 的方案数
• 换一个方向研究这个组合意义，尝试把 $j$ 给丢掉
• 先选出这 $k$ 个球，方案数 $\binom{n-ik}k$
• 而对于第一次选出的 $j$ 个球中的唯一一个染成 $m$ 的球，有两种情况
• （1）在第二次选出的 $k$ 个球内
• 这时我们可以在这 $k$ 个球中任选一个球染成 $m$ ，其他 $k-1$ 个球染成 $[1,m-1]$ ，其余 $n-ik-k$ 个球可以染成 $[1,m]$ （因为其余 $n-k$ 个球可以在第一次被选出也可以不被选出）
• 方案数为
• $k(m-1)^{k-1}m^{n-ik-k}$
• （2）在第一次选出而第二次未选出的 $j-k$ 个球内
• 这时候第二次选出的所有 $k$ 个球必须染成 $[1,m-1]$ ，对于剩下的 $n-ik-k$ 个球我们必须选出其中一个球，使得这个球被染成 $m$ ，并且在第一次被选出，这样剩下的 $n-ik-k-1$ 个球就可以任意染成 $[1,m]$ 内的颜色了
• 方案数为
• $(n-ik-k)(m-1)^km^{n-ik-k-1}$
• 于是
• $\sum_{j=k}^nj(m-1)^{j-1}\binom jk\binom{n-ik}j$
• $=\binom{n-ik}k(k(m-1)^{k-1}m^{n-ik-k}+(n-ik-k)(m-1)^km^{n-ik-k-1})$
• 于是答案也就是
• $ans=nm^n-m\sum_{i=0}^{n-1}\sum_{k=0}^{\lfloor\frac n{i+1}\rfloor}\frac {(-1)^k}{n-ik}\binom{n-ik}k(k(m-1)^{k-1}m^{n-ik-k}+(n-ik-k)(m-1)^km^{n-ik-k-1})$
• 预处理阶乘、逆元、阶乘逆元、 $m-1$ 的幂和 $m$ 的幂
• 复杂度 $O(n\log n)$
• orz zzq

Code

#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>

const int N = 3e5 + 5;

int n, m, ZZQ, fac[N], inv[N], invf[N], pm1[N], pm[N], ans;

int qpow(int a, int b)
{
	int res = 1;
	while (b)
	{
		if (b & 1) res = 1ll * res * a % ZZQ;
		a = 1ll * a * a % ZZQ;
		b >>= 1;
	}
	return res;
}

int C(int n, int m)
{
	return 1ll * fac[n] * invf[m] % ZZQ * invf[n - m] % ZZQ;
}

int main()
{
	//freopen("sequence.in", "r", stdin);
	//freopen("sequence.out", "w", stdout);
	
	std::cin >> n >> m >> ZZQ;
	fac[0] = inv[1] = invf[0] = invf[1] = 1;
	for (int i = 1; i <= n; i++)
		fac[i] = 1ll * fac[i - 1] * i % ZZQ;
	for (int i = 2; i <= n; i++)
		inv[i] = 1ll * (ZZQ - ZZQ / i) * inv[ZZQ % i] % ZZQ;
	for (int i = 2; i <= n; i++)
		invf[i] = 1ll * inv[i] * invf[i - 1] % ZZQ;
	pm1[0] = pm[0] = 1;
	for (int i = 1; i <= n; i++)
	{
		pm1[i] = 1ll * (m - 1) * pm1[i - 1] % ZZQ;
		pm[i] = 1ll * m * pm[i - 1] % ZZQ;
	}
	for (int i = 0; i < n; i++)
		for (int k = 0; i * k + k <= n; k++)
		{
			int delta = 0;
			if (k) delta = (1ll * k * pm1[k - 1] % ZZQ *
				pm[n - i * k - k] + delta) % ZZQ;
			if (n - i * k - k) delta = (1ll * (n - i * k - k)
				* pm[n - i * k - k - 1] % ZZQ * pm1[k] + delta) % ZZQ;
			delta = 1ll * delta * C(n - i * k, k) % ZZQ * inv[n - i * k] % ZZQ;
			if (k & 1) ans = (ans - delta + ZZQ) % ZZQ;
			else ans = (ans + delta) % ZZQ;
		}
	std::cout << (1ll * n * qpow(m, n) % ZZQ -
	1ll * m * ans % ZZQ + ZZQ) % ZZQ << std::endl;
	return 0;
}