Find a subset with given average

Problem

An integer array is given and an average value is given. Find a subset of the array which has the given average. For example input array is 1,3,4,5 and average is 2.5 then the subset {1,4} is the answer as its average is 2.5

Solution

We use dynamic programming to solve this problem. Suppose L is a 3-d boolean array. L[i,j,k]=true if using array element up to index i, a sum of j is possible using k number of elements from that range, and false otherwise. If we don't take the ith element then L[i,j,k]=L[i-1,j,k]. If we take the ith element then L[i,j,k]=L[i-1,j-arr[i],k-1]. Based on this formula we fill the 3-d array from bottom up. So Given the array, sum and count this array can tell whether a sum can be attained using k elements. If an average is given we do this operation n times where n is the count of the elements in the array. So if average is a, we search for L[n,a,1], L[n,2*a,2], ..., L[n,n*a,n]. If anyone of these searches returns true then we can obtain the given average. The function finds whether that sum is attainable with given count and returns the subset.

To return the subset we backtrack the 3-d array. If L[i,j,k]=L[i-1,j,k] then ith element should not be in the subset. If L[i,j,k]=L[i-1,j-arr[i],k-1] then that element is present in the subset. 

The complexity of the solution is O(n^4). O(n^3) for filling the 3-d array and then doing this process n times for finding average.

Code

import java.util.ArrayList;
import java.util.List;

public class IsAveragePossible
{
 public static void main(String[] args)
 {
  int[] arr =
  { 1, 4, 5, 3, 8 };
  double average = 3.5;
  for (int i = 1; i <= arr.length; ++i)
  {
   double sum = i * average;
   if (sum == (int) sum)
   {
    List < integer > result = isSumPossible(arr, (int) sum, i);
    if (result != null)
    {
     for (int num : result)
      System.out.print(num + ", ");
     System.out.println();
     return;
    }
   }
  }

  System.out.println("Not possible");
 }

 private static List < integer > isSumPossible(int[] arr, int sum, int count)
 {
  boolean[][][] memo = new boolean[arr.length + 1][sum + 1][count + 1];
  memo[0][0][0] = true;

  for (int i = 1; i <= arr.length; ++i)
   for (int j = 1; j <= sum; ++j)
    for (int k = 1; k <= count; ++k)
    {
     if (k == 1 && j == arr[i - 1])
      memo[i][j][k] = true;
     else if (arr[i - 1] > j)
      memo[i][j][k] = memo[i - 1][j][k];
     else
      memo[i][j][k] = memo[i - 1][j][k]
        || memo[i - 1][j - arr[i - 1]][k - 1];
    }
  if (memo[arr.length][sum][count] == false)
   return null;
  int i = arr.length;
  int j = sum;
  int k = count;
  List < integer > list = new ArrayList < integer >();
  while (i > 0 && j > 0 && k > 0)
  {
   if (memo[i][j][k] == memo[i - 1][j][k])
   {

   } else if (k == 1 && j == arr[i - 1])
   {
    list.add(arr[i - 1]);
    break;
   } else if (arr[i - 1] <= j
     && memo[i][j][k] == memo[i - 1][j - arr[i - 1]][k - 1])
   {
    j = j - arr[i - 1];
    k = k - 1;
    list.add(arr[i - 1]);
   }
   i--;
  }
  return list;
 }
}