2302. Count Subarrays With Score Less Than K

HardArrayBinary SearchPrefix SumSliding Window
Leetcode Link

Problem Description

The problem provides us with an array of positive integers, nums, and a number k. Our task is to find the number of contiguous non-empty subarrays within nums where the score of the subarray is less than k. The score of a subarray is calculated by multiplying the sum of the elements in the subarray by the length of the subarray. A subarray is just a sequence of elements from the array that are adjacent to each other.

Intuition

To solve this problem, we need to find a way to efficiently calculate the score for all possible subarrays in nums and count how many of them have a score less than k. Doing this directly would require examining every subarray separately, which can be very inefficient for large arrays.

The intuition behind the solution is to use a sliding window approach that allows us to calculate the score of subarrays dynamically as we expand and contract the window. By keeping track of the sum of elements currently in the window, we can determine the score quickly for the current window size.

Here's the general idea:

  • We start with a window at the beginning of the array.
  • We expand the window to the right by adding elements until the score exceeds or equals k.
  • When the score is equal to or greater than k, we shrink the window from the left by removing elements until the score is less than k again.
  • For each window size, we count the number of valid subarrays that can be formed, which is equivalent to the number of elements we can add to the right of the current window while maintaining a score less than k.

This approach ensures we only calculate the score for relevant subarrays and prevents unnecessary recalculations, leading us to the solution in an efficient manner.

Learn more about Binary Search, Prefix Sum and Sliding Window patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece๏ผš

What's the output of running the following function using input [30, 20, 10, 100, 33, 12]?

1def fun(arr: List[int]) -> List[int]:
2    import heapq
3    heapq.heapify(arr)
4    res = []
5    for i in range(3):
6        res.append(heapq.heappop(arr))
7    return res
8
1public static int[] fun(int[] arr) {
2    int[] res = new int[3];
3    PriorityQueue<Integer> heap = new PriorityQueue<>();
4    for (int i = 0; i < arr.length; i++) {
5        heap.add(arr[i]);
6    }
7    for (int i = 0; i < 3; i++) {
8        res[i] = heap.poll();
9    }
10    return res;
11}
12
1class HeapItem {
2    constructor(item, priority = item) {
3        this.item = item;
4        this.priority = priority;
5    }
6}
7
8class MinHeap {
9    constructor() {
10        this.heap = [];
11    }
12
13    push(node) {
14        // insert the new node at the end of the heap array
15        this.heap.push(node);
16        // find the correct position for the new node
17        this.bubble_up();
18    }
19
20    bubble_up() {
21        let index = this.heap.length - 1;
22
23        while (index > 0) {
24            const element = this.heap[index];
25            const parentIndex = Math.floor((index - 1) / 2);
26            const parent = this.heap[parentIndex];
27
28            if (parent.priority <= element.priority) break;
29            // if the parent is bigger than the child then swap the parent and child
30            this.heap[index] = parent;
31            this.heap[parentIndex] = element;
32            index = parentIndex;
33        }
34    }
35
36    pop() {
37        const min = this.heap[0];
38        this.heap[0] = this.heap[this.size() - 1];
39        this.heap.pop();
40        this.bubble_down();
41        return min;
42    }
43
44    bubble_down() {
45        let index = 0;
46        let min = index;
47        const n = this.heap.length;
48
49        while (index < n) {
50            const left = 2 * index + 1;
51            const right = left + 1;
52
53            if (left < n && this.heap[left].priority < this.heap[min].priority) {
54                min = left;
55            }
56            if (right < n && this.heap[right].priority < this.heap[min].priority) {
57                min = right;
58            }
59            if (min === index) break;
60            [this.heap[min], this.heap[index]] = [this.heap[index], this.heap[min]];
61            index = min;
62        }
63    }
64
65    peek() {
66        return this.heap[0];
67    }
68
69    size() {
70        return this.heap.length;
71    }
72}
73
74function fun(arr) {
75    const heap = new MinHeap();
76    for (const x of arr) {
77        heap.push(new HeapItem(x));
78    }
79    const res = [];
80    for (let i = 0; i < 3; i++) {
81        res.push(heap.pop().item);
82    }
83    return res;
84}
85

Solution Approach

The implementation of the solution uses a two-pointer (or sliding window) approach that keeps track of the current subarray being considered. These two pointers are denoted as i (the right pointer) and j (the left pointer), which represent the current bounds of the subarray.

Here is a step-by-step explanation of the algorithm:

  1. Initialize ans to 0; this will count the number of valid subarrays. Also, initialize s to 0; this will hold the sum of the elements of the current subarray. Initialize the left pointer j to 0, which represents the start of the current subarray.

  2. Iterate over each element v in nums using its index i, which acts as the end of the subarray. This loop will expand the window to the right.

    • Add the value of the current element, v, to s, which maintains the sum of the subarray from index j to i.

  3. While the score of the current subarray is greater than or equal to k (i.e., s * (i - j + 1) >= k), remove elements from the start of the subarray to reduce the score.

    • Subtract nums[j] from s to reduce the sum. This corresponds to "removing" the element at the start of the subarray.
    • Increment j to effectively shrink the window from the left.

  4. At this point, the sum multiplied by the window length is less than k. Therefore, for the current end of the window (i), we can count the number of valid subarrays that end at i as i - j + 1, since we can form a valid subarray by starting from any element between j and i.

    • Add i - j + 1 to ans, which accumulates the number of valid subarrays.

  5. After the iteration is complete, ans will hold the total number of valid subarrays, and we return this value as the final answer.

By using this algorithm, we avoid explicitly calculating the score for every possible subarray, and we efficiently count the valid subarrays by maintaining an ongoing sum and adjusting the window size. The algorithm has an overall time complexity of O(n) because each element is added to s and removed from s at most once, keeping the number of operations linear with the size of the input array.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Which algorithm should you use to find a node that is close to the root of the tree?

Example Walkthrough

Let's walk through the provided solution with an example. Suppose our array nums is [2, 1, 4, 1] and the number k is 8.

We will follow the steps described in the solution approach:

  1. We initialize ans to 0, which will store the total count of valid subarrays, s to 0 for the sum of the current subarray elements, and the left pointer j to 0.

  2. We begin iterating over nums with the right pointer i. For each element v in nums, we perform the following steps:

    • i = 0, v = 2: Add 2 to s. s becomes 2. (s * (i - j + 1)), which is (2 * (0 - 0 + 1)) = 2, is less than 8, so we add i - j + 1 = 0 - 0 + 1 = 1 to ans.

    • i = 1, v = 1: Add 1 to s. s becomes 3. (s * (i - j + 1)), which is (3 * (1 - 0 + 1)) = 6, is less than 8, so we add i - j + 1 = 1 - 0 + 1 = 2 to ans.

    • i = 2, v = 4: Add 4 to s. s becomes 7. (s * (i - j + 1)), which is (7 * (2 - 0 + 1)) = 21, is greater than 8, so we start shrinking the window from the left:

      • We subtract num[j] which is 2 from s and increment j. Now s is 5 and j is 1. The updated score is (5 * (2 - 1 + 1)) = 10, which is still greater than 8.
      • We again subtract num[j], now 1, from s and increment j. Now s is 4 and j is 2. The score is (4 * (2 - 2 + 1)) = 4, which is less than 8. Now we add i - j + 1 = 2 - 2 + 1 = 1 to ans.

    • i = 3, v = 1: Add 1 to s. s becomes 5. (s * (i - j + 1)), which is (5 * (3 - 2 + 1)) = 10, is again greater than 8. We need to shrink the window:

      • We do not need to remove any elements from the window since j is already at 2 and score 10 is from subarray [4, 1]. Now since we cannot shrink the window and the score is greater than k, we simply move on.

At the end of the iteration, we have the total number of valid subarrays which is the value of ans. By adding the counts at each step, ans = 1 + 2 + 1 = 4. Thus, [2], [2, 1], [1] and [1, 4] are valid subarrays whose scores are less than 8, and the final answer is 4.

Solution Implementation

1from typing import List
2
3class Solution:
4    def countSubarrays(self, nums: List[int], k: int) -> int:
5        # Initialize answer, current sum, and start index of the window
6        answer = current_sum = start_index = 0
7      
8        # Iterate through the array with index and value
9        for end_index, value in enumerate(nums):
10            # Add the current number to the current sum
11            current_sum += value
12          
13            # While the product of the sum of the current subarray
14            # and the length of the subarray is at least k,
15            # shrink the window from the left (increase start_index)
16            while current_sum * (end_index - start_index + 1) >= k:
17                current_sum -= nums[start_index]
18                start_index += 1
19          
20            # The number of subarrays ending with the current number
21            # is the length of the current window (end_index - start_index + 1)
22            answer += end_index - start_index + 1
23      
24        # Return the total number of subarrays found
25        return answer
26
1class Solution {
2    public long countSubarrays(int[] nums, long k) {
3        long count = 0; // To store the number of subarrays
4        long sum = 0; // Sum of the elements in the current subarray
5        int start = 0; // Start index for the current subarray
6
7        // Traverse through the array starting from the 0th element
8        for (int end = 0; end < nums.length; ++end) {
9            sum += nums[end]; // Add the current element to sum
10          
11            // Shrink the subarray from the left if the condition is violated
12            // sum * length should be less than k
13            while (sum * (end - start + 1) >= k) {
14                sum -= nums[start]; // Removing the element from the start of subarray
15                start++; // Increment the start index
16            }
17          
18            // At this point, for each element nums[end], we find how many subarrays ending at 'end' are valid
19            // The number of valid subarrays is given by the difference btw current end and the new start position
20            count += end - start + 1;
21        }
22        return count; // Return the total count of valid subarrays
23    }
24}
25
1class Solution {
2public:
3    long long countSubarrays(vector<int>& nums, long long k) {
4        long long count = 0; // Initialize the count of subarrays
5        long long sum = 0;   // Initialize the sum of elements in the current subarray
6        int start = 0;       // Start index for our sliding window
7
8        // Iterate through each element in the array
9        for (int end = 0; end < nums.size(); ++end) {
10            sum += nums[end]; // Add the current element to the sum
11
12            // If the constraint is violated (average * number of elements >= k)
13            // increment start index to reduce the number of elements and sum
14            while (sum * (end - start + 1) >= k) {
15                sum -= nums[start++]; // Remove the element pointed by start from sum and increment start
16            }
17
18            // Count the number of subarrays ending with nums[end]. If no elements were removed in the while loop,
19            // this will include all subarrays starting from nums[0..end]. If some elements were removed,
20            // it will include all subarrays starting from nums[start..end].
21            count += end - start + 1;
22        }
23
24        return count; // Return the total count of valid subarrays
25    }
26};
27
1function countSubarrays(nums: number[], k: number): number {
2    let count = 0; // Initialize the count of subarrays
3    let sum = 0;   // Initialize the sum of the elements in the current subarray
4    let start = 0; // Start index for our sliding window
5
6    // Iterate through each element in the array
7    for (let end = 0; end < nums.length; end++) {
8        sum += nums[end]; // Add the current element to the sum
9
10        // While the average of the subarray multiplied by the number of elements
11        // is greater than or equal to k, adjust the start index to shrink the subarray
12        while (sum * (end - start + 1) >= k) {
13            sum -= nums[start]; // Remove the element pointed to by start from the sum
14            start++; // Increment start index to reduce the subarray size
15        }
16
17        // Count the number of subarrays ending with nums[end]. This includes all
18        // subarrays starting from nums[start..end]. If the while loop above didn't
19        // run, it would also include all subarrays from nums[0..end].
20        count += end - start + 1;
21    }
22
23    return count; // Return the total count of valid subarrays
24}
25
Not Sure What to Study? Take the 2-min Quiz๏ผš

Which algorithm should you use to find a node that is close to the root of the tree?

Time and Space Complexity

Time Complexity

The given code uses a sliding window technique to count the number of subarrays whose elements product is less than k. In this code, two pointers (i and j) are used, which move forward through the array without stepping backwards. This results in each element being considered only once by each pointer, resulting in a linear traversal.

Thus, the time complexity of this algorithm is O(n), where n is the number of elements in the array nums. This is because both pointers i and j can only move from the start to the end of the array once, and the operations inside the for-loop and while-loop are all constant time operations.

Space Complexity

The space complexity is determined by the extra space used aside from the input. In this case, only a fixed number of variables (ans, s, j, i, v) are used. These do not depend on the size of the input array. Therefore, the space complexity of the code is O(1), which is constant space complexity since no additional space that grows with the input size is used.

Learn more about how to find time and space complexity quickly using problem constraints.

Fast Track Your Learning with Our Quick Skills Quiz:

Which two pointer technique does Quick Sort use?

Recommended Readings

Binary Search Speedrun

Binary Search Speedrun For each of the Speedrun questions you will be given a binary search related problem and a corresponding multiple choice question The multiple choice questions are related to the techniques and template s introduced in the binary search section It's recommended that you have gone through at

Prefix Sum

Prefix Sum The prefix sum is an incredibly powerful and straightforward technique Its primary goal is to allow for constant time range sum queries on an array What is Prefix Sum The prefix sum of an array at index i is the sum of all numbers from index 0 to i By

Sliding Window Maximum

https algomonster s3 us east 2 amazonaws com cover_photos stack svg Sliding Window Maximum Monotonic Stack We have an array and a sliding window defined by a start index and an end index The sliding window moves from left of the array to right There are always k elements in

Got a question?ย Ask the Teaching Assistantย anything you don't understand.

Still not clear? Ask in the Forum, ย Discordย orย Submitย the part you don't understand to our editors.

โ†
โ†‘TA ๐Ÿ‘จโ€๐Ÿซ