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:
- Start with the range of window as [i,j] = [0,K-1]. Then iterate over remaining elements.
- 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)).
- Keep track of the max average at every step
- 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