Count the Number of Arrays with K Matching Adjacent Elements

Count the Number of Arrays with K Matching Adjacent Elements

You are given three integers n, m, k. A good array arr of size n is defined as follows:

• Each element in arr is in the inclusive range [1, m].
• Exactly k indices i (where 1 <= i < n) satisfy the condition arr[i - 1] == arr[i].

Return the number of good arrays that can be formed.

Since the answer may be very large, return it modulo 109 + 7.

 

Example 1:

Input: n = 3, m = 2, k = 1

Output: 4

Explanation:

• There are 4 good arrays. They are [1, 1, 2], [1, 2, 2], [2, 1, 1] and [2, 2, 1].
• Hence, the answer is 4.

Example 2:

Input: n = 4, m = 2, k = 2

Output: 6

Explanation:

• The good arrays are [1, 1, 1, 2], [1, 1, 2, 2], [1, 2, 2, 2], [2, 1, 1, 1], [2, 2, 1, 1] and [2, 2, 2, 1].
• Hence, the answer is 6.

Example 3:

Input: n = 5, m = 2, k = 0

Output: 2

Explanation:

• The good arrays are [1, 2, 1, 2, 1] and [2, 1, 2, 1, 2]. Hence, the answer is 2.

 

Constraints:

• 1 <= n <= 105
• 1 <= m <= 105
• 0 <= k <= n - 1

 

解题思路

  半年前做出来的题结果现在不会了（）。看到数据范围其实可以基本确定是组合数学的题了，但就是想不到怎么做。dp 的话状态至少是二维，也做不了，当时赛时是发现转移方程很像杨辉三角然后猜过的。

  由于一开始想不出组合数学的解法，还是先给出 dp 的状态定义。定义 $f(i,j)$ 表示前 $i$ 个数中有 $j$ 对相邻元素相同的方案数，根据第 $i$ 个数是否与第 $i-1$ 个数相同进行状态划分，状态转移方程为 $f(i,j) = (m-1)f(i-1,j) + f(i-1,j-1)$。其中初始状态 $f(1,0)=m$。显然整个 dp 的复杂度为 $O(nk)$。

  注意到转移方程与杨辉三角中第 $i$ 行第 $j$ 列元素的表达式 $C_{i}^{j} = C_{i-1}^{j} + C_{i-1}^{j-1}$ 非常相似，因此我们可以以 dp 的状态转移方程来模拟杨辉三角，看看是否存在规律。前 $5$ 行的结果如下，其中行和列都从 $0$ 开始计数。

$$\begin{array}{|c|c|c|c|c|} \hline m & 0 & 0 & 0 & 0 \\
\hline m(m-1) & m & 0 & 0 & 0 \\
\hline m(m-1)^2 & 2m(m-1) & m & 0 & 0 \\
\hline m(m-1)^3 & 3m(m-1)^2 & 3m(m-1) & m & 0 \\
\hline m(m-1)^4 & 4m(m-1)^3 & 6m(m-1)^2 & 4m(m-1) & m \\ \hline \end{array}$$

  根据上述的结果，容易发现规律第 $i$ 行第 $j$ 列的结果是 $C_{i}^{j} \cdot m(m-1)^{i-j}$，对应 $f(i+1,j)$。由于最终的答案为 $f(n,k)$，因此对应的结果就是 $C_{n-1}^{k} \cdot m(m-1)^{n-1-k}$。

  下面给出组合数学的解法。对于一个长度为 $n$ 的序列，共有 $n-1$ 对相邻元素，现在我们要使得其中的 $k$ 对相邻元素相同，意味着有 $n-1-k$ 对相邻元素不同。参考 2 2 2 3 3 4 4，一共有 $6$ 对相邻数字，其中有 $4$ 对相同。另外不同的 $2$ 对我们可以看作是 $2$ 条分割线，把序列分成了 $3$ 部分，且每一个部分内的数字都相同。对于更一般的情况，当我们要选出 $k$ 对相同的相邻元素，等价于选出 $n-1-k$ 对不同的相邻元素将数组分成 $n-k$ 个部分，且相邻两部分内的数字不同（同部分的相同）。对应的方案数就是 $C_{n-1}^{n-1-k} \cdot m (m-1)^{n-1-k}$（其中第一个部分有 $m$ 种选择），可以发现与上面推到的结果一致。

  AC 代码如下，时间复杂度为 $O(k)$：

class Solution {
public:
    int mod = 1e9 + 7;
    
    int countGoodArrays(int n, int m, int k) {
        auto qmi = [&](int a, int k) {
            int ret = 1;
            while (k) {
                if (k & 1) ret = 1ll * ret * a % mod;
                a = 1ll * a * a % mod;
                k >>= 1;
            }
            return ret;
        };
        auto C = [&](int a, int b) {
            int u = 1, d = 1;
            for (int i = 1, j = a; i <= b; i++, j--) {
                u = 1ll * u * j % mod;
                d = 1ll * d * i % mod;
            }
            return 1ll * u * qmi(d, mod - 2) % mod;
        };
        return 1ll * C(n - 1, k) * m % mod * qmi(m - 1, n - 1 - k) % mod;
    }
};

  记录这道题的另外一个目的是当我重新写这道题时，写出的状态转移方程与上述的形式并不相同，而是 $f(i,j) = (m-1)\sum\limits_{u=0}^{j}{f(i-u-1, j-u)}$，即按照最后一段有多少对相同的相邻元素进行划分。实际上该转移方程与 $f(i,j) = (m-1)f(i-1,j) + f(i-1,j-1)$ 完全是等价的，和完全背包的优化方式一样，可以用 $f(i-1,j)$ 来代替 $(m-1)\sum\limits_{u=1}^{j}{f(i-u-1, j-u)}$ 的部分。然而我更希望的是能一次性写出优化后的形式，但我目前还没想到很好的方法。不过还是能总结出一些规律的，如果写出了类似这样的式子，又发现划分的数量本身没有限制（不是状态的大小），比如完全背包中物品可以无限选择，本题中最后一段相同的相邻元素数量没有限制，那么就可以根据不选或者只选一个来划分状态，而不是枚举选多少个。而像多重背包（物品选择数量有限）就不可以用这种方法了。相关的题目还有 D1. The Endspeaker (Easy Version)，当时就是因为这个问题没写出来。

 

参考资料

  枚举右维护左 组合数学【力扣周赛 430】：https://www.bilibili.com/video/BV13f68YjE7o/