【LeetCode】每日一题 —— No.3405

LeetCode 3405 统计恰好有 K 个相等相邻元素的数组数目（DP 构造型）

题目概述

我们需要统计长度为 n 的数组 arr 满足如下条件的方案数：

1. 每个元素在区间 [1, m] 之间
2. 恰好存在 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. 与前一个元素相等：

   • 只可选择 1 种值（即与 arr[i-1] 相同）
   • 所以 dp[i][j] += dp[i-1][j-1] * 1

2. 与前一个元素不同：

   • 可选择 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;
    }
};

测试用例建议