好数组的总数 | 豆包MarsCode AI 刷题

问题理解

1. 输入：

   • N：数组 A 的长度。
   • K：数组 B 中每个元素的取值范围为 1 到 K。
   • A：一个长度为 N 的数组，每个元素在 1 到 K 的范围内。

2. 输出：

   • 满足条件的好数组的总数。

3. 条件：

   • 数组 B 的每个元素 B[i] 必须在 1 到 K 的范围内。
   • 对于所有从 0 到 N-2 的索引 i，满足 B[i] + B[i+1] = A[i] + A[i+1]。

解题思路

1. 动态规划：

   • 我们可以使用动态规划来解决这个问题。定义一个二维数组 dp[i][j]，表示在第 i 个位置上，B[i] 取值为 j 时，满足条件的好数组的数量。
   • 初始条件：dp[0][j] = 1 对于所有 1 <= j <= K，因为第一个位置可以取任意值。
   • 状态转移：对于每个位置 i，我们需要计算 dp[i][j]，其中 j 是 B[i] 的取值。根据条件 B[i] + B[i+1] = A[i] + A[i+1]，我们可以推导出 B[i+1] 的取值范围。

2. 状态转移方程：

   • 对于每个位置 i 和 B[i] 的取值 j，我们需要找到所有可能的 B[i+1] 的取值 k，使得 j + k = A[i] + A[i+1]。
   • 因此，dp[i+1][k] 可以由 dp[i][j] 转移而来。

3. 最终结果：

   • 最终结果是所有 dp[N-1][j] 的和，其中 1 <= j <= K。

public class Main {
    public static int solution(int N, int K, int[] A) {
        // 初始化 dp 数组
        int[][] dp = new int[N][K + 1];
        
        // 初始条件
        for (int j = 1; j <= K; j++) {
            dp[0][j] = 1;
        }
        
        // 状态转移
        for (int i = 0; i < N - 1; i++) {
            for (int j = 1; j <= K; j++) {
                for (int k = 1; k <= K; k++) {
                    if (j + k == A[i] + A[i + 1]) {
                        dp[i + 1][k] += dp[i][j];
                    }
                }
            }
        }
        
        // 计算最终结果
        int result = 0;
        for (int j = 1; j <= K; j++) {
            result += dp[N - 1][j];
        }
        
        return result;
    }

    public static void main(String[] args) {
        System.out.println(solution(3, 2, new int[] { 1, 2, 1 }) == 2);
        System.out.println(solution(4, 3, new int[] { 1, 3, 2, 1 }) == 1);
        System.out.println(solution(2, 1, new int[] { 1, 1 }) == 1);
    }
}

代码解释

1. 初始化：

   • dp[0][j] = 1 表示第一个位置可以取任意值。

2. 状态转移：

   • 对于每个位置 i 和 B[i] 的取值 j，我们找到所有可能的 B[i+1] 的取值 k，使得 j + k = A[i] + A[i+1]。

3. 最终结果：

   • 最终结果是所有 dp[N-1][j] 的和。

通过这个思路和伪代码框架，你可以逐步实现代码并验证其正确性。

关键思路提示

1. 初始化：

   • 确保 dp[0][j] 初始化为 1，因为第一个位置可以取任意值。

2. 状态转移：

   • 在状态转移过程中，确保 k 的取值在 1 到 K 的范围内，避免数组越界。

3. 边界条件：

   • 确保在状态转移时，i 的范围是从 0 到 N-2，因为我们需要计算到倒数第二个位置。

4. 结果计算：

   • 最终结果是所有 dp[N-1][j] 的和，确保累加所有可能的 j 值。