【AGC002F】Leftmost Ball

题目

题目链接：https://atcoder.jp/contests/agc002/tasks/agc002_f
给你 (n) 种颜色的球，每个球有 (k) 个，把这 (n imes k) 个球排成一排，把每一种颜色的最左边出现的球涂成白色(初始球不包含白色)，求有多少种不同的颜色序列，答案对 (10^9+7) 取模。
(n,kleq 2000)。

思路

最后肯定是有 (n) 个白球，以及 (n) 种颜色的球各 (k-1) 个。观察到一个序列合法，当且仅当任意前缀，白球的数量都不小于其他颜色球的颜色数。
设 (f[i][j]) 表示放了 (i) 个白球，以及 (j) 种颜色的球的方案数。注意可以任意放，不一定是放在一个前缀内。
考虑目前从左开始第一个出现的空位放什么：

• 如果放白球，那么从 (f[i-1][j]) 转移过来，条件是 (i>j)。
• 如果放新的颜色的球，那么从 (f[i][j-1]) 转移过来，且这一个空位必须填这个球，这种颜色剩余的 (k-2) 个球，就需要在剩余的 (nk-i-(j-1)(k-1)-1) 个空位内随意放。

所以转移为

[f[i][j]=f[i-1][j]+f[i][j-1] imes (n-j+1) imes inom{nk-i-(j-1)(k-1)-1}{k-2} ]

特判一下 (k=1) 的情况即可。
时间复杂度 (O(nk))。

代码

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

const int N=2010,MOD=1e9+7;
int n,k;
ll fac[N*N],inv[N*N],f[N][N];

ll fpow(ll x,ll k)
{
	ll ans=1;
	for (;k;k>>=1,x=x*x%MOD)
		if (k&1) ans=ans*x%MOD;
	return ans;
}

void init()
{
	fac[0]=inv[0]=1;
	for (int i=1;i<=n*k;i++) fac[i]=fac[i-1]*i%MOD;
	inv[n*k]=fpow(fac[n*k],MOD-2);
	for (int i=n*k-1;i>=1;i--) inv[i]=inv[i+1]*(i+1)%MOD;
}

ll C(int n,int m)
{
	if (n<m) return 0;
	return fac[n]*inv[m]%MOD*inv[n-m]%MOD;
}

int main()
{
	scanf("%d%d",&n,&k);
	if (k==1) return printf("1"),0;
	init();
	for (int i=1;i<=n;i++)
	{
		f[i][0]=1;
		for (int j=1;j<=i;j++)
		{
			f[i][j]=f[i][j-1]*(n-j+1)%MOD*C(n*k-i-(j-1)*(k-1)-1,k-2)%MOD;
			if (j<i) f[i][j]=(f[i][j]+f[i-1][j])%MOD;
		}
	}
	cout<<f[n][n];
	return 0;
}