Question

I got asked this question on a interview for Google a couple of weeks ago, I didn’t quite get the answer and I was wondering if anyone here could help me out.

You have an array with n elements. The elements are either 0 or 1.
You want to split the array into k contiguous subarrays. The size of each subarray can vary between ceil(n/2k) and floor(3n/2k). You can assume that k << n.
After you split the array into k subarrays. One element of each subarray will be randomly selected.

Devise an algorithm for maximizing the sum of the randomly selected elements from the k subarrays.
Basically means that we will want to split the array in such way such that the sum of all the expected values for the elements selected from each subarray is maximum.

You can assume that n is a power of 2.

Example:

Array: [0,0,1,1,0,0,1,1,0,1,1,0]
n = 12
k = 3
Size of subarrays can be: 2,3,4,5,6

Possible subarrays [0,0,1] [1,0,0,1] [1,0,1,1,0]
Expected Value of the sum of the elements randomly selected from the subarrays: 1/3 + 2/4 + 3/5 = 43/30 ~ 1.4333333 

Optimal split: [0,0,1,1,0,0][1,1][0,1,1,0]
Expected value of optimal split: 1/3 + 1 + 1/2 = 11/6 ~ 1.83333333

Answer 1

I don’t know if this is still an open question or not, but it seems like the OP has managed to add enough clarifications that this should be straightforward to solve. At any rate, if I am understanding what you are saying this seems like a fair thing to ask in an interview environment for a software development position.

Here is the basic O(n^2 * k) solution, which should be adequate for small k (as the interviewer specified):

def best_val(arr, K):
  n = len(arr)
  psum = [ 0.0 ]
  for x in arr:
    psum.append(psum[-1] + x)
  tab = [ -100000 for i in range(n) ]
  tab.append(0)
  for k in range(K):
    for s in range(n - (k+1) * ceil(n/(2*K))):
      terms = range(s + ceil(n/(2*K)), min(s + floor((3*n)/(2*K)) + 1, n+1))
      tab[s] = max( [ (psum[t] - psum[s]) / (t - s) + tab[t] for t in terms ])
  return tab[0]

I used the numpy ceil/floor functions but you basically get the idea. The only `tricks’ in this version is that it does windowing to reduce the memory overhead to just O(n) instead of O(n * k), and that it precalculates the partial sums to make computing the expected value for a box a constant time operation (thus saving a factor of O(n) from the inner loop).