Problem of The Day: K Radius Subarray Averages

Problem Statement

leetcode problem link

Intuition

The problem requires calculating the average of a subarray centered at each index, where the subarray has a total length of 2k + 1. For indices where this window goes out of bounds, we return -1. The first idea is to use a sliding window approach to efficiently calculate the sum of each subarray in constant time, rather than recomputing the sum for every window from scratch.

Approach

We initialize a result array res with all elements set to -1, since not every index will be able to have a valid average computed (especially near the start and end).

We use a sliding window of size 2k + 1 and maintain a running sum curr_sum. As we iterate through the array using a pointer end, we add nums[end] to curr_sum. Once we’ve reached the window size, we compute the average for the middle index of the window and store it in res[middle_index]. Then we slide the window forward by subtracting the element at start and incrementing both start and middle_index.

Complexity

• Time complexity:
  \(O(n)\)
  Each element is added and removed from the sliding window sum exactly once.

• Space complexity:
  \(O(n)\)
  We use an output array res of the same size as the input array.

Code

class Solution:
    def getAverages(self, nums: List[int], k: int) -> List[int]:
        N = len(nums)
        res = [-1] * N
        window_size = (k * 2) + 1
        curr_sum = 0
        start = 0
        middle_index = k
        for end in range(N):
            curr_sum += nums[end]
            if end >= window_size - 1:
                res[middle_index] = math.floor(curr_sum / window_size)
                curr_sum -= nums[start]
                start += 1
                middle_index += 1
        return res

Editorial

Approach 1: Prefix Sum

class Solution:
    def getAverages(self, nums: List[int], k: int) -> List[int]:
        # When a single element is considered then its average will be the number itself only.
        if k == 0:
            return nums

        window_size = 2 * k + 1
        n = len(nums)
        averages = [-1] * n

        # Any index will not have 'k' elements in it's left and right.
        if window_size > n:
            return averages

        # Generate 'prefix' array for 'nums'.
        # 'prefix[i + 1]' will be sum of all elements of 'nums' from index '0' to 'i'.
        prefix = [0] * (n + 1)
        for i in range(n):
            prefix[i + 1] = prefix[i] + nums[i]

        # We iterate only on those indices which have atleast 'k' elements in their left and right.
        # i.e. indices from 'k' to 'n - k'
        for i in range(k, n - k):
            leftBound, rightBound = i - k, i + k
            subArraySum = orefix[rightBound + 1] - prefix[leftBound]
            average = subArraySum // window_size
            averages[i] = average

        return averages

Approach 2: Sliding Window

class Solution:
    def getAverages(self, nums: List[int], k: int) -> List[int]:
        averages = [-1] * len(nums)
        # When a single element is considered then its average will be the number itself only.
        if k == 0:
            return nums

        window_size = 2 * k + 1
        n = len(nums)

        # Any index will not have 'k' elements in it's left and right.
        if window_size > n:
            return averages

        # First get the sum of first window of the 'nums' arrray.
        window_sum = sum(nums[:window_size])
        averages[k] = window_sum // window_size

        # Iterate on rest indices which have at least 'k' elements
        # on its left and right sides.
        for i in range(window_size, n):
            # We remove the discarded element and add the new element to get current window sum.
            # 'i' is the index of new inserted element, and
            # 'i - (window size)' is the index of the last removed element.
            window_sum = window_sum - nums[i - window_size] + nums[i]
            averages[i - k] = window_sum // window_size

        return averages