[LeetCode] 2090. K Radius Subarray Averages

問題

https://leetcode.com/problems/k-radius-subarray-averages/description/

解答と解説

class Solution:
    def getAverages(self, nums: List[int], k: int) -> List[int]:
        prefixSum = [nums[0]]
        for i in range(1, len(nums)):
            prefixSum.append(prefixSum[-1] + nums[i])
        
        res = [-1 for _ in range(len(nums))]
        for i in range(k, len(nums) - k):
            start, end = i - k, i + k
            total = prefixSum[end] - prefixSum[start - 1] if start >= 1 else prefixSum[end]
            res[i] = total // (2 * k + 1)
            
        return res

numsの要素数をnとした時、

• Time Complexity: O(n)
• Space Complexity: O(n)

-1であるかどうか考えるのは面倒なので最初に初期化しておく。

更新されるindexはkからlen(nums)-k-1までの要素であり、平均値はPrefixSumを使うことでO(n)で計算できる。