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 | 合法輸出 | 測試性能邊界 |
