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

问题理解

1. 输入：

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

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，因为第一个位置上的任何取值都是合法的。
   • 对于每个位置 i，遍历所有可能的 B[i] 取值 j，并根据 B[i] + B[i+1] = A[i] + A[i+1] 的条件，更新 dp[i+1][k]。

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] 的累加值，即所有可能的 B[N-1] 取值的 dp 值之和

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++) {
                if (dp[i][j] > 0) {
                    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]。
   • 更新 dp[i+1][k] 为 dp[i][j] 的累加值。

3. 最终结果：

   • 最终结果为 dp[N-1][j] 的累加值，即所有可能的 B[N-1] 取值的 dp 值之和。

通过这种方式，我们可以有效地计算出满足条件的好数组的总数。