LeetCode 805 Split Array With Same Average (推荐 预处理 类01背包dp 详解 640ms -＞ 16ms)

You are given an integer array nums.

You should move each element of nums into one of the two arrays A and B such that A and B are non-empty, and average(A) == average(B).

Return true if it is possible to achieve that and false otherwise.

Note that for an array arr, average(arr) is the sum of all the elements of arr over the length of arr.

Example 1:

Input: nums = [1,2,3,4,5,6,7,8]
Output: true
Explanation: We can split the array into [1,4,5,8] and [2,3,6,7], and both of them have an average of 4.5.

Example 2:

Input: nums = [3,1]
Output: false

Constraints:

• 1 <= nums.length <= 30
• 0 <= nums[i] <= 104

题目链接：https://leetcode.com/problems/split-array-with-same-average/

题目大意：问能否把数组分成两部分使得两部分的平均值相等

题目分析：设分成两部分的个数和总和分别为n1, sum1, n2, sum2，则有sum1/n1 = sum2/n2等价于sum1/n1 = (sum - sum1)/(n - n1)，化简得sum1/n1 = sum/n ，于是问题转化为是否能在n个数里找到n1个数，使它们的均值等于总均值，因sum1，n1是整数，故可过滤出可能的n1，又因要分成两部分，必然有min(n1, n2) <= n/2，可借此减少枚举数量，进一步将问题转化为是否能在n个数里找n1个数，使它们的和为C（C = n1 * sum / n），对于每个枚举的n1，C可当做定值，此问题可用01背包模型做，dp[i][j][k]表示枚举到i，选了j个数是否能合成k，转移方程：

dp[i][j][k] |= dp[i - 1][j - 1][k - A[i]]
dp[i][j][k] |= dp[i - 1][j][k]

641ms，时间击败40.71%

class Solution {
    public boolean splitArraySameAverage(int[] A) {
        int n = A.length;
        if (n == 1 || (n == 2 && (A[0] != A[1]))) {
            return false;
        }
        int sum = 0;
        for (int i = 0; i < n; i++) {
            sum += A[i];
        }
        
        int[] candidates = new int[n >> 1];
        int cnt = 0;
        for (int i = 1; i <= n / 2; i++) {
            if ((i * sum) % n == 0) {
                candidates[cnt++] = i;
            }
        }
        
        if (cnt == 0) {
            return false;
        }

        boolean[][][] dp = new boolean[n][n / 2 + 1][sum];
        dp[0][0][0] = true;
        dp[0][1][A[0]] = true;
        for (int i = 1; i < n; i++) {
            for (int j = 0; j <= Math.min(i + 1, n / 2); j++) {
                for (int k = 0; k < sum; k++) {
                    if (j > 0 && k >= A[i]) {
                        dp[i][j][k] |= dp[i - 1][j - 1][k - A[i]];   
                    }
                    dp[i][j][k] |= dp[i - 1][j][k];
                    // if (dp[i][j][k]) {
                    //     System.out.println("dp["+i+"]["+j+"]["+k+"] = " + dp[i][j][k]);
                    // }
                }
            }
        }
        
        for (int i = 0; i < cnt; i++) {
            int tN = candidates[i];
            int tSum = (tN * sum) / n;
            for (int j = 0; j < n; j++) {
                if (dp[j][tN][tSum]){
                    return true;
                }
            }
        }

        return false;
    }
}

01背包有个空间二维到一维的优化在此也适用，因每个数字A[i]只能选一次，只需将二层循环逆序即可，这样可以保证用来更新当前i阶段dp[j][k]的值必然来自i-1阶段，这里可能不太好理解，可以反过来看为什么正序枚举不对，如若正序枚举，假设当前对A[i]进行决策（选或不选），且j枚举到2，如果选A[i] => dp[2][k] |= dp[1][k-A[i]]，j继续枚举到3，继续选A[i] => dp[3][k'] |= dp[2][k'- A[i]]，若k = k' - A[i]，即j枚举到3时选了A[i]的状态是从j枚举到2时选了A[i]的状态推导出来的，也就是说此时A[i]被选了两次，不符合题意，这就容易理解为什么j逆序枚举是正确的了，j枚举到3时，要用j等2时的dp值来更新，当前i阶段j等于2时的dp值还未被推导，故使用的必然是i-1阶段的值

188ms，时间击败61.79%

class Solution {
    public boolean splitArraySameAverage(int[] A) {
        int n = A.length;
        if (n == 1 || (n == 2 && (A[0] != A[1]))) {
            return false;
        }
        int sum = 0;
        for (int i = 0; i < n; i++) {
            sum += A[i];
        }
        
        int[] candidates = new int[n >> 1];
        int cnt = 0;
        for (int i = 1; i <= n / 2; i++) {
            if ((i * sum) % n == 0) {
                candidates[cnt++] = i;
            }
        }
        
        if (cnt == 0) {
            return false;
        }

        boolean[][] dp = new boolean[n / 2 + 1][sum];
        dp[0][0] = true;
        dp[1][A[0]] = true;
        for (int i = 1; i < n; i++) {
            for (int j = Math.min(i + 1, n / 2); j > 0; j--) {
                for (int k = A[i]; k < sum; k++) {
                    dp[j][k] |= dp[j - 1][k - A[i]];   
                }
            }
        }
        
        for (int i = 0; i < cnt; i++) {
            int tN = candidates[i];
            int tSum = (tN * sum) / n;
            if (dp[tN][tSum]){
                return true;
            }
        }

        return false;
    }
}

还可以进一步优化，当前的主要耗时点在于第三层的枚举其实存在大量无用的值（dp值为false的k值），如果某个阶段的dp值为false，那么之后就不可能再变成true（这是显然的），于是可以考虑直接存能组成的数字，将dp[i]变成一个hashset，表示选i个数字能组成的集合

16ms，时间击败97.17%

class Solution {
    public boolean splitArraySameAverage(int[] A) {
        int n = A.length;
        if (n == 1 || (n == 2 && (A[0] != A[1]))) {
            return false;
        }
        int sum = 0;
        for (int i = 0; i < n; i++) {
            sum += A[i];
        }
        
        boolean ok = false;
        for (int i = 1; i <= n / 2; i++) {
            if ((i * sum) % n == 0) {
                ok = true;
                break;
            }
        }
        
        if (!ok) {
            return false;
        }

        int m =  n / 2 + 1;
        Set<Integer>[] dp = new HashSet[m];
        for (int i = 0; i < m; i++) {
            dp[i] = new HashSet<>();
        }
        dp[0].add(0);
        for (int i = 0; i < n; i++) {
            for (int j = Math.min(i + 1, n / 2); j > 0; j--) {
                for (int num : dp[j - 1]) {
                    dp[j].add(num + A[i]);
                }
            }
        }

        for (int i = 1; i < m; i++) {
            if ((i * sum) % n == 0 && dp[i].contains(i * sum / n)) {
                return true;
            }
        }

        return false;
    }
}