Search code examples
arraysalgorithmdata-structuressub-array

Possibly simpler O(n) solution to find the Sub-array of length K (or more) with the maximum average


I saw this question on a coding competition site.

Suppose you are given an array of n integers and an integer k (n<= 10^5, 1<=k<=n). How to find the sub-array(contiguous) with maximum average whose length is more than k.

There's an O(n) solution presented in research papers(arxiv.org/abs/cs/0207026.), linked in a duplicate SO question. I'm posting this as a separate question since I think I have a similar method with a simpler explanation. Do you think there's anything wrong with my logic in the solution below?

Here's the logic:

  1. Start with the range of window as [i,j] = [0,K-1]. Then iterate over remaining elements.
  2. For every next element, j, update the prefix sum**. Now we have a choice - we can use the full range [i,j] or discard the range [i:j-k] and keep [j-k+1:j] (i.e keep the latest K elements). Choose the range with the higher average (use prefix sum to do this in O(1)).
  3. Keep track of the max average at every step
  4. Return the max avg at the end

** I calculate the prefix sum as I iterate over the array. The prefix sum at i is the cumulative sum of the first i elements in the array.

Code:

def findMaxAverage(nums, k):
    prefix = [0]
    for i in range(k):
        prefix.append(float(prefix[-1] + nums[i]))
    mavg = prefix[-1]/k
    lbound = -1
    for i in range(k,len(nums)):
        prefix.append(prefix[-1] + nums[i])
        cavg = (prefix[i+1] - prefix[lbound+1])/(i-lbound)
        altavg = (prefix[i+1] - prefix[i-k+1])/k
        if altavg > cavg:
            lbound = i-k
            cavg = altavg
        mavg = max(mavg, cavg)

    return mavg

Solution

  • Consider k = 3 and sequence

    3,0,0,2,0,1,3
    

    Output of your program is 1.3333333333333333. It has found subsequence 0,1,3, but the best possible subsequence is 2,0,1,3 with average 1.5.