当前位置：首页 > web >正文

leetcode-3405 统计恰好有k个相等相邻数组的个数

web 2025/6/19 21:19:33

题目描述

给你三个整数 n ，m ，k 。长度为 n 的好数组 arr 定义如下：

arr 中每个元素都在闭区间 [1, m] 中。
恰好有k个下标 i（其中 1 <= i < n）满足 arr[i - 1] == arr[i] 。

请你返回可以构造出的好数组数目。

由于答案可能会很大，请你将它对 10^9 + 7 取余后返回。

解题过程

阶段一

一开始我把它当成了一个动态规划的问题，如何理解呢？

假定容纳n个元素的数组，可选值范围为m，需要k个满足题目条件的子序列的数组个数有dp(n,m,k)个。

我们可以从n-1的情况递推到n的情况，因为无非就是多了一个容纳元素的格子，这个格子可以是与前者相同的元素，这种元素我们称之为A类元素，它可以构造题目需要的k个子序列中的一个。另一种情况就是它不选择A类元素，而是选择普通元素，即不会创造相等邻接子序列的元素，我们称之为X元素，除了数组的第一个元素(必定为X元素，因为第一个元素前面没有任何元素可以接)有m种取值以外，其他的X元素都只有m-1种取值，因为它不能与前面的元素相同(否则就变成了A类元素)。

所以当我们开始递推时：

当新加入的格子选择A类元素时，不同的数组个数有 dp(n - 1,m,k - 1) 个，因为A类元素只有一种取值。

当新加入的格子选择X类元素时，不同的数组个数有dp(n - 1, m , k) * (m - 1)个

所以递推公式为

dp(n,k) = dp(n-1,k-1) + dp(n - 1, k) * (m - 1)

这里因为m始终不变，所以省略不写了，本质上是一个二维dp。

当k = 0时，dp(n,0) 有 m * pow(m - 1, n -1)个不同的数组，这由排列组合的知识得到。

因此便可以按顺序填值了。

于是有了第一种代码：

int countGoodArrays(int n, int m, int k) {
	int search[n][k + 1] = {};
	unsigned long long resolves = m;for(int i = 0;i < n; i++){
		search[i][0] = resolves;
		resolves = (resolves * (m - 1)) % K;}for(int i = 1;i < n;i ++)for(int q = 1;q < k + 1;q ++){
		search[i][q] = (search[i-1][q] * (m - 1) + search[i-1][q-1]) % K;}
	return search[n-1][k];
}

阶段二

该代码的内存是可以优化的：

因为dp数组的前面每一行用完了后面就用不到了，所以可以删掉，所以引入了轮转的优化代码：

int countGoodArrays2(int n, int m, int k) {
	vector<int> tmp(k + 1, 0),ret;unsigned long long resolves = m;
	tmp[0] = resolves;	for(int i = 1;i < n; i++){
		resolves = (resolves * (m - 1)) % K;
		ret.emplace_back(resolves);for(int q = 0;q < k;q ++)
		ret.emplace_back(((unsigned long long)tmp[q + 1] * (m - 1) + tmp[q]) % K);
		tmp.clear();
		ret.swap(tmp);
		ret.reserve(k + 1);}return tmp[k];
}

通过每一次swap来使得最终的dp数组占用只有两行。

阶段三

后面提测发现题目给的n,m,k都会很大，这就导致时间和空间基本怎么样都会超。这里怀疑方法有问题了，又仔细地想了一下。

新思路：

给定了大小为n的数组，除了数组第一个元素必须是X元素以外，剩下的可以是A类元素可以是X类元素。那么这里可以用组合数C来从n - 1个位置里面选择出k个位置填充A类元素，剩下的都是X类元素，但是关键是组合数需要计算阶乘，这里数值计算会很大！所以相当于本题是在考察大数下的组合数计算。注意到题目有说：

由于答案可能会很大，请你将它对 10^9 + 7 取余后返回。

这里才知道要用费马小定理结合模逆元来计算组合数。

在模运算中除法可以改成逆元的乘法，这是模运算的性质。

由于题目给的MOD是个素数，所以可以用费马小定理来计算一个数的逆元。

当p是素数时，a的逆元可以转换成a的高阶幂次，所以要做幂运算，我们通过快速幂运算来完成这一点：

long long qpow(long long a, long long b) {long long res = 1;
    a %= K;while (b > 0) {if (b & 1) res = res * a % K;
        a = a * a % K;
        b >>= 1;}return res;
}

这样我们就能求出每个数的逆元，而由于输入n > m，所以最高需要求的阶乘就是n!，而n最大是1e5，所以我们通过数组来保存计算的阶乘。

for(int i = 0; i <= MX; i ++){if(i == 0){
			pw.emplace_back(1);
			continue;}
		pw.emplace_back((unsigned long long)pw[pw.size() - 1] * i % K);}

然后通过逆元的递推公式：

得到对应的每个阶乘的逆元数组：

anv[MX] = qpow(pw[MX],K-2);for(int i = MX - 1;i >= 0; i--){
		anv[i] = (i + 1) * (unsigned long long)anv[i + 1] % K;}

最后代入即可求得组合数。

组合数出来之后，对于每个A类元素，取值只有一种可能，对于第一个X类元素，取值有m种可能，后面的每个X类元素有m - 1种可能，乘起来即可，这里涉及到幂次，所以又可以用qpow快速幂函数。

x = (x * qpow(m-1,n - k -1)) % K;

最终的代码为：

#include <vector>
vector<int> anv,pw;
class Solution {
public:unsigned long long K = 1000000000 + 7;int MX = 1e5;
long long qpow(long long a, long long b) {long long res = 1;
    a %= K;while (b > 0) {if (b & 1) res = res * a % K;
        a = a * a % K;
        b >>= 1;}return res;
}int C(int m, int n)
{if(!n || m == n) return 1;if(n > m) return 0;return (unsigned long long)pw[m] * anv[n] % K * anv[m - n] % K;
}
int countGoodArrays(int n, int m, int k) {if(pw.empty()){
        pw.reserve(MX + 1);
        anv.reserve(MX + 1);for(int i = 0; i <= MX; i ++){if(i == 0){
			pw.emplace_back(1);continue;}
		pw.emplace_back((unsigned long long)pw[pw.size() - 1] * i % K);}
	anv[MX] = qpow(pw[MX],K-2);for(int i = MX - 1;i >= 0; i--){
		anv[i] = (i + 1) * (unsigned long long)anv[i + 1] % K;}}int ret = 0;unsigned long long resolves = C(n - 1, k);unsigned long long x = m;
	x = (x * qpow(m-1,n - k -1)) % K;
	ret = (resolves * x) % K;return ret;
}
};