「AtCoder AGC002F」Leftmost Ball

Description

现有 (n) 种颜色（不是白色）的球，每种各 (k) 个。将这些球排列好后，对于每种颜色，将这种颜色的最前面的球涂成白色。试问最终得到的颜色序列有多少种。

Hint

• (1le n, kle 2 imes 10^3)

Solution

我们换一种理解方式理解题目：现有白色球 (n) 个，其他不同颜色的球各 (k - 1) 个，要求将这些球摆放，要求得到的颜色序列的任意一个前缀中白色球的个数不小于其他颜色的颜色种类数。

如果去掉最后那个限制那就简单了，遗憾的是白球的特殊性使得它必须被 特殊考虑。

---

首先基本可以断定这个题是 组合计数动态规划。

最可能想到的是，设 (f(i)) 表示前 (i) 个的方案数。然而这样的话状态信息是严重缺失的，但因为数据规模的缘故难以再添加维度。

从位置方面难以突破，因此 从颜色开始考虑。

---

又可以想到设 (f(i)) 为放置颜色 (0sim i) 的方案数（(0) 表示白色），一种颜色一次性一起放。显然我们还没有考虑到白色的特殊性，于是不妨加一维——设 (f(i, j)) 为 放了 (i) 个白球， (j) 种其他颜色的球 的方案数，白球一个个放，其他的一次性放。注意，这个状态在 (i ge j) 时才有意义。

---

状态转移情况分析：

• 这个状态可以由 (f(i - 1, j)) 转移过来：放置一个白球到达 (f(i, j)) 的状态。
• 也可以从 (f(i, j - 1)) 转移过来：放置 (k-1) 个一种颜色的球。这样就不能直接加了，情况比较复杂。
• • 首先毋庸置疑的是第一个位置必须放。
  • 然后在剩下的 (nk - i - 1 - (k - 1)(j - 1)) 个空位中随便放余下的 (k - 2) 个即可。
  • 这里就有 (egin{pmatrix} nk - i - 1 - (k - 1)(j - 1) \ k - 2 end{pmatrix}) 种方案。最后因为颜色也是可以任意选的，因此还要乘上剩下的颜色数。

不难得出状态转移方程：

[f(i, j) = f(i - 1, j) + f(i, j - 1) imes egin{pmatrix} nk - i - 1 - (k - 1)(j - 1) \ k - 2 end{pmatrix} imes (n - j + 1) ]

时空复杂度 (O(n^2))。

Code

/*
 * Author : _Wallace_
 * Source : https://www.cnblogs.com/-Wallace-/
 * Problem : AtCoder AGC002F Leftmost Ball
 */
#include <iostream>

using namespace std;
const int N = 2e3 + 5;
const int mod = 1e9 + 7;
typedef long long LL;

LL fastpow(LL a, LL b) {
    if (!b) return 1;
    LL t = fastpow(a, b / 2);
    if (b & 1) return t * t % mod * a % mod;
    else return t * t % mod;
}

LL fact[N * N], invf[N * N];
void prework(int n) {
    fact[0] = 1ll;
    for (int i = 1; i <= n; i++)
        fact[i] = fact[i - 1] * i % mod;
    invf[n] = fastpow(fact[n], mod - 2ll);
    for (int i = n - 1; ~i; i--)
        invf[i] = invf[i + 1] * (i + 1) % mod;
}

LL comb(int n, int m) {
    return fact[n]
         * invf[m] % mod
         * invf[n - m] % mod;
}

int k, n;
LL f[N][N];

signed main() {
    prework(4e6);
    cin >> n >> k;

    if (k == 1) {
        cout << 1 << endl;
        return 0;
    }

    for (int i = 0; i <= n; i++)
        f[i][0] = 1ll;
    
    for (int i = 1; i <= n; i++)
        for (int j = 1; j <= i; j++) {
            f[i][j] += (f[i][j - 1] * (n - j + 1) % mod
                     * comb(n - i + (n - j + 1) * (k - 1) - 1, k - 2) % mod
                     + f[i - 1][j]) % mod;
        }
    
    cout << f[n][n] << endl;
    return 0;
}