Solution -「UOJ #450」复读机

$\mathcal{Description}$

  Link.

  求从 $m$ 种颜色，每种颜色无限多的小球里选 $n$ 个构成排列，使得每种颜色出现次数为 $d$ 的倍数的排列方案数，对 $19491001$ 取模。

  $n\le10^9$，

• $m\le10^3$，$d=3$；

• $m\le5\times10^5$，$d\le2$。

$\mathcal{Solution}$

  分 $d=1,2,3$ 求解。

  当 $d=1$，每个位置 $m$ 种方案，答案为 $m^n$。

  当 $d=2$，偶数序列的 EGF 为 $G(x)=\sum_{i=0}^{+\infty}\frac{x^{2i}}{(2i)!}=\frac{e^x+e^{-x}}2$，那么答案为：

$$\begin{aligned} n![x^n]G^m(x)&=n![x^n]\left( \frac{e^x+e^{-x}}2 \right)^m\\ &=\frac{n!}{2^m}[x^n]\left( \sum_{i=0}^m\binom{m}ie^{(2i-m)x} \right)\\ &=2^{-m}\sum_{i=0}^m\binom{m}i(2i-m)^n \end{aligned}$$

  第二步到第三步用到常见的 $e^{ax}=\sum_{i=0}^{+\infty}\frac{a^i}{i!}x^i$。此时就能 $\mathcal O(m\log n)$ 求出答案了。

  当 $d=3$，$3$ 的倍数数的 EGF 为 $G(x)=\sum_{i=0}^{+\infty}[3|i]\frac{x^i}{i!}$，这个不太好算，来一发单位根反演：

$$\begin{aligned} G(x)&=\sum_{i=0}^{+\infty}[3|i]\frac{x^i}{i!}\\ &=\frac{1}3\sum_{i=0}^{+\infty}\frac{x^i}{i!}\sum_{j=0}^2\omega_3^{ij}\\ &=\frac{1}3\sum_{j=0}^2\sum_{i=0}^{+\infty}\frac{(\omega_3^jx)^i}{i!}\\ &=\frac{1}3\sum_{j=0}^2e^{\omega_3^jx} \end{aligned}$$

  接着求答案，暴力展开三项式幂：

$$\begin{aligned} n![x^n]G^m(x)&=\frac{n!}{3^m}[x^n]\left( \sum_{j=0}^2e^{\omega_3^jx} \right)^m\\ &=\frac{n!}{3^m}[x^n]\left( \sum_{a+b+c=m}\binom{m}{a,b,c}e^{(a\omega_3^0+b\omega_3^1+c\omega_3^2)x} \right)\\ &=3^{-m}\sum_{a+b+c=m}\binom{m}{a,b,c}(a\omega_3^0+b\omega_3^1+c\omega_3^2)^n \end{aligned}$$

  注意到 $c=m-a-b$，所以多重组合数的值就是 $\frac{m!}{a!b!c!}$，该式能在 $\mathcal O(m^2\log n)$ 的时间内算出。实际上该式就是 $d=2$ 的情况的扩展，由于 $\omega_2^{0,1}=\pm1$，所以亦能从该式推回 $d=2$ 的情况。

$\mathcal{Code}$

/* Clearink */

#include <cstdio>

#define rep( i, l, r ) for ( int i = l, rpbound##i = r; i <= rpbound##i; ++i )
#define per( i, r, l ) for ( int i = r, rpbound##i = l; i >= rpbound##i; --i )

const int MOD = 19491001, MAXM = 5e5, INV2 = MOD + 1 >> 1, INV3 = 12994001;
const int W[] = { 1, 663067, 18827933 };
int n, m, d, fac[MAXM + 5], ifac[MAXM + 5];

inline int mul ( long long a, const int b ) { return a * b % MOD; }
inline int sub ( int a, const int b ) { return ( a -= b ) < 0 ? a + MOD : a; }
inline int add ( int a, const int b ) { return ( a += b ) < MOD ? a : a - MOD; }
inline int mpow ( int a, int b ) {
	int ret = 1;
	for ( ; b; a = mul ( a, a ), b >>= 1 ) ret = mul ( ret, b & 1 ? a : 1 );
	return ret;
}

inline void init () {
	fac[0] = 1;
	rep ( i, 1, m ) fac[i] = mul ( i, fac[i - 1] );
	ifac[m] = mpow ( fac[m], MOD - 2 );
	per ( i, m - 1, 0 ) ifac[i] = mul ( i + 1, ifac[i + 1] );
}

inline int comb ( const int n, const int m ) {
	return n < m ? 0 : mul ( fac[n], mul ( ifac[m], ifac[n - m] ) );
}

int main () {
	scanf ( "%d %d %d", &n, &m, &d ), init ();
	if ( d == 1 ) return printf ( "%d\n", mpow ( m, n ) ), 0;
	if ( d == 2 ) {
		int ans = 0;
		rep ( i, 0, m ) {
			ans = add ( ans, mul ( comb ( m, i ), mpow ( sub ( i << 1, m ), n ) ) );
		}
		printf ( "%d\n", mul ( ans, mpow ( INV2, m ) ) );
		return 0;
	}
	// $d is now smaller than 1000.
	int ans = 0;
	rep ( i, 0, m ) rep ( j, 0, m - i ) {
		int k = m - i - j;
		ans = add ( ans, mul ( mul ( ifac[i], mul ( ifac[j], ifac[k] ) ),
			mpow ( add (
				mul ( i, W[0] ), add ( mul ( j, W[1] ), mul ( k, W[2] ) ) ), n ) ) );
	}
	printf ( "%d\n", mul ( ans, mul ( fac[m], mpow ( INV3, m ) ) ) );
	return 0;
}