LeetCode 1343. Number of Sub-arrays of Size K and Average Greater than or Equal to Threshold

, , ,

Question

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
2
3
4
5
6
7
8
9
10
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).

Input: arr = [11,13,17,23,29,31,7,5,2,3], k = 3, threshold = 5
Output: 6
Explanation: The first 6 sub-arrays of size 3 have averages greater than 5. Note that averages are not integers.

Input: arr = [7,7,7,7,7,7,7], k = 7, threshold = 7
Output: 1

Solution

It’s sliding window problem, and we just need to maintain the k size window and calculate the average. Also we can use one variable to store the sum that we add so far to reduce the duplicate calculations. In addtion, you can use prefix sum array to optimize.

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
// time:O(n - k + 1 ~ n) space:O(1)
    public int numOfSubarrays(int[] arr, int k, int threshold) {
        // prefix sum
        if (arr == null || arr.length == 0) return 0;
        int sum = 0, res = 0;
        for (int i = 0; i < arr.length; i++) {
            sum += arr[i];
            if (i >= k - 1) {
                if (sum >= threshold * k) res++;
                sum -= arr[i - k + 1];
            }
        }
        return res;
    }

    // time:O(n) space:O(n)
    public int numOfSubarrays2(int[] arr, int k, int threshold) {
        if (arr == null || arr.length == 0) return 0;
        int n = arr.length;
        int[] prefix = new int[n + 1];
        for (int i = 1; i <= n; i++) {
            prefix[i] = prefix[i - 1] + arr[i - 1];
        }
        int index = 0;
        int res = 0;
        while (index < n && index + k <= n) {
            int sum = prefix[index + k] - prefix[index];
            if (sum >= k * threshold) res++;
            index++;
        }
        return res;
    }