805. Split Array With Same Avera

Description

In a given integer array A, we must move every element of A to either list B or list C. (B and C initially start empty.)

Return true if and only if after such a move, it is possible that the average value of B is equal to the average value of C, and B and C are both non-empty.

Example :
Input:
[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 the average of 4.5.

Note:

• The length of A will be in the range [1, 30].
• A[i] will be in the range of [0, 10000].

Solution

HashMap + DP (TLE)

类似01背包问题，注意从后向前遍历！

class Solution {
    public boolean splitArraySameAverage(int[] arr) {
        int n = arr.length;
        int totalSum = getSum(arr);
        
        Map<Integer, Set<Integer>> sum2Count = new HashMap<>();
        Set<Integer> initial = new HashSet<>();
        initial.add(0);
        sum2Count.put(0, initial);
        
        for (int i = 0; i < n; ++i) {
            for (int sum = totalSum; sum >= arr[i]; --sum) {    // back to front!
                if (!sum2Count.containsKey(sum - arr[i])) {
                    continue;
                }
                
                if (!sum2Count.containsKey(sum)) {
                    sum2Count.put(sum, new HashSet<>());
                }
                Set<Integer> tmp = new HashSet<>(); // add to tmp to avoid ConcurrentModityException

                for (int count : sum2Count.get(sum - arr[i])) {
                    if (count + 1 != n) {
                        tmp.add(count + 1);
                    }
                }
                
                sum2Count.get(sum).addAll(tmp);
                for (int count : sum2Count.get(sum)) {
                    if (isAvgEqual(sum, count, totalSum, n)) {
                        System.out.println(sum2Count.entrySet());
                        System.out.println(sum + "/" + count + " and " + totalSum + "/" + n);
                        return true;
                    }
                }
            }
        }
        
        return false;
    }
    // both of the count should larger than 0!
    private boolean isAvgEqual(int sum1, int count1, int sum2, int count2) {
        return count1 > 0 && count2 > 0 && count1 * sum2 == count2 * sum1;
    }
    
    private int getSum(int[] arr) {
        int sum = 0;
        for (int v : arr) {
            sum += v;
        }
        return sum;
    }
}