【LeetCode】每日一题 —— No.3405
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 | 合法输出 | 测试性能边界 |