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;
}