美团2024届秋招笔试第二场编程真题第六题：背包问题的前缀和优化

题目出自牛客

暴力

容易想到暴力方法如下

数组含义定义： dp[i][j] 表示在填到第i个数时，总和如果为j的构造方式。
初始化： 那么i = 0时， 除了j= nums[0]和j= 0. 其他的构造方式都为1.
递推： 那么为了方便考虑，就枚举i>0时， 如果当前填写数字[1,sum]会对当前的数组dp[i]造成的影响。
同时考虑到这是个markov过程，所以使用滚动数组优化。
代码如下：

import java.util.Scanner;

// 注意类名必须为 Main, 不要有任何 package xxx 信息
public class Main {
    final static int MOD = 1000_000_007;
    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        int n = in.nextInt();
        int[] nums = new int[n]; int sum = 0;
        for (int i = 0; i < n; i++) {
            nums[i] = in.nextInt();
            sum += nums[i];
        }

        // markov 过程 ， 滚动数组简化空间
        // dp[i][j]: i表示填到第几个, J 表示当前总和。
        int[][] dp = new int[2][sum+1];
        for (int i = 0; i <= sum; i++) dp[0][i] = 1; 
        dp[0][0] = 0; dp[0][nums[0]] = 0;
        for (int i = 1; i < n; i++) {
            int index = i % 2;
            int pre = 1-index;

            for (int kk = 0; kk <= sum; kk++) dp[index][kk] = 0; // 滚动数组清零
            for (int j = 1; j <= sum; j ++) { // 当前填的数字
                if (j == nums[i]) continue; // 不同
                for (int count = j+1; count <= sum; count ++) { // 当前的组合值。
                    int preN = count - j;  // 填
                    dp[index][count] += dp[pre][preN];
                    dp[index][count] %= MOD;
                }
            }
        }

        int index = (n-1)%2;
        System.out.println(dp[index][sum]);
    }
}

前缀和优化

总感觉如果把二维数组表示成在坐标i处，当前和<=j的可行构造数能把复杂度降很多。
拿起笔推了一下，容易发现如下公式，

讨论过程：

因此可以把数组dp的含义修改为 在坐标i处，当前总和小于j的可行构造数。 这样可以将内层的一个for循环编程前缀和之间的相减。 代码如下。

import java.util.Scanner;

// 注意类名必须为 Main, 不要有任何 package xxx 信息
public class Main {
    final static int MOD = 1000_000_007;
    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        int n = in.nextInt();
        int[] nums = new int[n]; int sum = 0;
        for (int i = 0; i < n; i++) {
            nums[i] = in.nextInt();
            sum += nums[i];
        }

        long[][] dp = new long[n][sum+1];
        for (int i = 1; i <= sum; i++) {
            if (i != nums[0]) {
                dp[0][i] = 1 + dp[0][i-1]; 
            } else {
                dp[0][i] = dp[0][i-1]; 
            }
        }
        for (int i = 1; i < n; i++) {
            for (int j = 1; j <= sum; j ++) { // 当前总和
                // 如果不考虑不能重复， 和为j的出现次数为 dp[i-1][j-1].
                // 考虑的话就需要 i-1 处sum不能 j-nums[i]
                // j - nums[i] > 0. 出现的次数 为 dp[i-1][j - nums[i]] - dp[i-1][j - nums[i]-1]
                dp[i][j] = (dp[i][j-1] + dp[i-1][j-1])%MOD;
                if (j-nums[i] > 0) {
                    dp[i][j] -= dp[i-1][j - nums[i]] - dp[i-1][j - nums[i]-1];
                }
                while (dp[i][j] < 0) dp[i][j] += MOD;
                dp[i][j] = dp[i][j] % MOD;
            }
        }
        long res = dp[n-1][sum]-dp[n-1][sum-1];
        while (res < 0) res += MOD;
        System.out.println(res);
    }
}

前后的运行时间如下：
优化前：

优化后