LeetCode 3405 统计恰好有 K 个相等相邻元素的数组数目(DP 构造型)
题目概述
我们需要统计长度为 n
的数组 arr
满足如下条件的方案数:
- 每个元素在区间
[1, m]
之间 - 恰好存在
k
个位置i (1 ≤ i < n)
满足arr[i] == arr[i - 1]
也就是说,整个数组中,恰好有 k
个相邻元素对相等。最终返回构造这些数组的方案数,对 1e9 + 7
取模。
示例说明
示例 1:
输入:n = 3, m = 2, k = 1
输出:4
合法数组为:[1,1,2]
, [1,2,2]
, [2,1,1]
, [2,2,1]
问题本质
本题本质上是一个构造型动态规划问题,并非从已有元素中选取,而是一步步构造合法的数组。
我们要统计的是:有多少种方法可以构造出一个长度为 n
的数组,包含恰好 k
个“相邻相等对”。
动态规划设计
1. 状态定义
设 dp[i][j]
表示**前 i
个元素中,恰好有 j
个“相邻相等对”**的方案数。
我们从左到右一个个构造数组元素。
i
表示数组长度(即前缀长度)j
表示当前构造中已有的相等对数
2. 初始状态
dp[1][0] = m
:第一个元素有m
种取值,没有相邻对。
3. 状态转移方程
我们从 dp[i-1][j]
推导出 dp[i][j]
:
新增元素的两种选择:
-
与前一个元素相等:
- 只可选择 1 种值(即与
arr[i-1]
相同) - 所以
dp[i][j] += dp[i-1][j-1] * 1
- 只可选择 1 种值(即与
-
与前一个元素不同:
- 可选择
m - 1
种不同的值 - 所以
dp[i][j] += dp[i-1][j] * (m - 1)
- 可选择
4. 完整转移公式
dp[i][j] = dp[i-1][j] * (m - 1) // 当前位和前一位不相等+ dp[i-1][j-1] * 1 // 当前位和前一位相等(构成一个新对)
注意对 j==0
和 j==i-1
边界的处理。
5. 目标值
我们要求的就是 dp[n][k]
。
代码实现
C++ 实现(TLE)
class Solution {
public:int countGoodArrays(int n, int m, int k) {const int MOD = 1e9 + 7;vector<int> prev(k + 1), curr(k + 1);prev[0] = m;for (int i = 2; i <= n; ++i) {for (int j = 0; j <= k && j < i; ++j) {long long same = (j > 0 ? prev[j - 1] : 0);long long diff = prev[j] * 1LL * (m - 1) % MOD;curr[j] = (same + diff) % MOD;}prev.swap(curr);}return prev[k];}
};
c++ 优化 1 (TLE)
常规的 O(n * k) 动态规划在极限数据下(n = 1e5)可能会超时。
优化版本:滚动数组 + 预处理 + 取模运算
- 优化点一:滚动数组 + 二维变一维
- 原始 dp[i][j] 是二维状态,但只依赖 dp[i-1][*]
- 可以使用一维滚动数组,空间压缩至 O(k)
- 优化点二:取模乘法处理要放在中间,避免溢出
状态定义
- dp[i][j]:长度为 i,且恰好有 j 个相邻相等对的数组数量。
状态转移
我们从 i=2 到 n,转移为:dp[i][j] = dp[i-1][j-1] + dp[i-1][j] * (m - 1)
- dp[i-1][j-1]:表示第 i 个数与第 i-1 个数相同,构成一个新对,只有 1 种取值(等于前一个)
- dp[i-1][j] * (m-1):当前数与前一个不同,有 m-1 种取值
初始状态
- dp[1][0] = m
第一个元素有 m 种选择,不可能形成相邻对
class Solution {
public:int countGoodArrays(int n, int m, int k) {const int MOD = 1e9 + 7;vector<int> dp(k + 1);dp[0] = m;for (int i = 2; i <= n; ++i) {vector<int> newDp(k + 1, 0);for (int j = 0; j <= k && j < i; ++j) {// arr[i] == arr[i-1],增加一个相等对if (j > 0) {newDp[j] = (newDp[j] + dp[j - 1]) % MOD;}// arr[i] != arr[i-1],保持已有对数,新增不同值newDp[j] = (newDp[j] + (long long)dp[j] * (m - 1)) % MOD;}dp = move(newDp);}return dp[k];}
};
c++优化2
使用 快速幂 + 预处理阶乘逆元 求组合数 C(n - 1, g - 1)
使用快速幂计算 (m - 1)^(g - 1)
class Solution {
public:static constexpr int MOD = 1e9 + 7;vector<long long> fact, inv_fact;long long fast_pow(long long base, long long exp) {long long res = 1;while (exp > 0) {if (exp & 1) res = res * base % MOD;base = base * base % MOD;exp >>= 1;}return res;}void init(int n) {fact.resize(n + 1);inv_fact.resize(n + 1);fact[0] = 1;for (int i = 1; i <= n; ++i) {fact[i] = fact[i - 1] * i % MOD;}inv_fact[n] = fast_pow(fact[n], MOD - 2);for (int i = n - 1; i >= 0; --i) {inv_fact[i] = inv_fact[i + 1] * (i + 1) % MOD;}}long long C(int n, int k) {if (k < 0 || k > n) return 0;return fact[n] * inv_fact[k] % MOD * inv_fact[n - k] % MOD;}int countGoodArrays(int n, int m, int k) {init(n);int g = n - k;long long comb = C(n - 1, g - 1);long long ways = comb * m % MOD;ways = ways * fast_pow(m - 1, g - 1) % MOD;return ways;}
};
测试用例建议
输入 | 输出 | 说明 |
---|---|---|
n=1, m=1, k=0 | 1 | 唯一合法数组 [1] |
n=2, m=1, k=1 | 1 | [1,1] |
n=2, m=1, k=0 | 0 | 无法构造不同值 |
n=3, m=2, k=1 | 4 | 题目示例 |
n=100000, m=100000, k=0 | 合法输出 | 测试性能边界 |