Leetcode - Number of Sub-arrays of Size K and Average Greater than or Equal to Threshold

violet posted @ May 05, 2020 06:08:47 AM in 算法 with tags Algorithm Golang PrefixSum , 226 阅读

https://leetcode.com/problems/number-of-sub-arrays-of-size-k-and-average-greater-than-or-equal-to-threshold/

Given an array of integers arr and two integers k and threshold.

Return the number of sub-arrays of size k and average greater than or equal to threshold.

Example 1:

Input: arr = [2,2,2,2,5,5,5,8], k = 3, threshold = 4
Output: 3
Explanation: Sub-arrays [2,5,5],[5,5,5] and [5,5,8] have averages 4, 5 and 6 respectively. All other sub-arrays of size 3 have averages less than 4 (the threshold).

Use prefix sum as a helper.

func numOfSubarrays(arr []int, k int, threshold int) int {
    sums := make([]int, len(arr))
    sums[0] = arr[0]
    for i := 1; i < len(arr); i++ {
        sums[i] = sums[i-1] + arr[i]
    }
    left := 0
    right := left
    result := 0
    for right < len(arr) {
        if right - left + 1 == k {
            var sum int
            if left > 0 {
                sum = sums[right]-sums[left-1]
            } else {
                sum = sums[right]
            }
            if sum/k >= threshold {
                result++
            }
            right++
            continue
        }
        if right - left + 1 < k {
            right++
        } else if right - left + 1 > k {
            left++
        }
    }
    return result
}