Equal Average Partition

Given an array A with non negative numbers, divide the array into two parts such that the average of both the parts is equal.

Return both parts (If exist). If there is no solution. return an empty list.

NOTE: If a solution exists, you should return a list of exactly 2 lists of integers A and B which follow the following condition :

numElements in A <= numElements in B

If numElements in A = numElements in B, then A is lexicographically smaller than B ( https://en.wikipedia.org/wiki/Lexicographical_order )

If multiple solutions exist, return the solution where length(A) is minimum. If there is still a tie, return the one where A is lexicographically smallest.

Array will contain only non negative numbers. **Input Format**

First andonly argument is an integer array A. **Output Format**

Return 2D array consisting of two rows where each row denoted a partition. **Example Input**

Input 1:

 A = [1 7 15 29 11 9]

Example Output

Output 1:

 [9 15] [1 7 11 29]

Example Explanation

Explanation 1:

 The average of part is (15+9)/2 = 12, average of second part elements is (1 + 7 + 11 + 29) / 4 = 12

package Questions;

import java.util.*;

class Solution {
    // mem[i][j][k]==0 -> unknown answer
    // mem[i][j][k]==1 -> answer True
    // mem[i][j][k]==2 -> answer False
    // mem[i][j][k] -> Whether it is possible to form a subset
    // of length -> j , sum -> k , from sub-array -> [i,length)
    byte[][][] mem;

    public ArrayList<ArrayList<Integer>> avgset(ArrayList<Integer> arrA) {
        int lenA = arrA.size();
        if (lenA <= 1)
            return new ArrayList<ArrayList<Integer>>();
        // Sorting the array because we want answer to be sorted as well
        Collections.sort(arrA);
        int sumA = 0;
        for (int x : arrA)
            sumA += x;
        mem = new byte[lenA][lenA / 2 + 1][sumA + 1];
        // We need to split A into B and C, the length of B can be [1, A.length / 2], we
        // consider them one by one:
        // B should have the same mean value as A, so sumB / lenOfB = sumA / lenOfA,
        // which is equivalent to sumB = sumA * lenOfB / lenOfA.
        // All elements here are integers, so sumB must be an integer, this gives
        // our first criteria (sumA * lenOfB) % A.length == 0.
        for (int lenB = 1; lenB <= arrA.size() / 2; ++lenB) {
            if (sumA * lenB % lenA != 0)
                continue;
            int needSum = sumA * lenB / lenA;
            if (isPossible(arrA, 0, lenB, needSum)) {
                ArrayList<ArrayList<Integer>> ans = new ArrayList<>();
                // List -> B
                ans.add(new ArrayList<Integer>());
                // List -> C
                ans.add(new ArrayList<Integer>());
                for (int j = 0; j < lenA; j++) {
                    int x = arrA.get(j);
                    if (ans.get(0).size() < lenB) {
                        // if answer is possible if we include x in arrB
                        // O(1) Time because of memoization
                        if (isPossible(arrA, j + 1, lenB - ans.get(0).size() - 1, needSum - x)) {
                            ans.get(0).add(x);
                            needSum -= x;
                        } else
                            ans.get(1).add(x);
                    } else {
                        ans.get(1).add(arrA.get(j));
                    }
                }
                return ans;
            }
        }
        return new ArrayList<ArrayList<Integer>>();
    }

    // To check whether we can form a subset of length -> lenB and sum = needSum
    public boolean isPossible(List<Integer> arrA, int fromIndex, int len, int sum) {
        if (sum == 0 && len == 0 && fromIndex <= arrA.size())
            return true;
        if (sum <= 0 || fromIndex >= arrA.size() || len > arrA.size() - fromIndex || len <= 0)
            return false;
        // If answer already exists
        if (mem[fromIndex][len][sum] != 0)
            return mem[fromIndex][len][sum] == 1;
        // 2 options
        // 1st -> Don't include the element at "index" in our needSum
        // 2nd -> Include the element at "index"
        if (isPossible(arrA, fromIndex + 1, len, sum)
                || isPossible(arrA, fromIndex + 1, len - 1, sum - arrA.get(fromIndex))) {
            mem[fromIndex][len][sum] = 1;
            return true;
        }
        mem[fromIndex][len][sum] = 2;
        return false;
    }
}