LeetCode 2302 - Count Subarrays With Score Less Than K

LeetCode Problem 2302

Difficulty: 🔴 Hard
Topics: Array, Binary Search, Sliding Window, Prefix Sum

Solution

Problem Understanding

The problem asks us to count the number of contiguous subarrays of a given array nums whose score is strictly less than a given threshold k. The score of a subarray is defined as the product of its sum and its length. Formally, for a subarray nums[i:j] with length len = j - i + 1 and sum sum(nums[i:j]), the score is sum(nums[i:j]) * len.

The input consists of an array of positive integers nums and a positive integer k. The expected output is a single integer representing the total count of subarrays that satisfy the score condition.

The constraints indicate that nums can have up to 10^5 elements, and individual elements as well as k can be very large (up to 10^5 for elements, up to 10^15 for k). This immediately rules out any brute-force solution that generates all subarrays explicitly, as the number of subarrays scales quadratically with n. Edge cases include arrays with all small values, arrays with large values, arrays where no subarray satisfies the condition, and arrays where all single-element subarrays satisfy the condition.

Approaches

A brute-force approach would iterate over all possible subarrays, compute the sum for each subarray, multiply it by its length, and check whether it is less than k. While this is straightforward and correct, it has a time complexity of O(n^2) for generating subarrays and O(n) for computing their sums if not optimized with prefix sums, resulting in a total time complexity of O(n^3) in the naive case. Using prefix sums reduces it to O(n^2), but this is still too slow for n = 10^5.

The key observation for a more optimal approach is that the score function is monotonic with respect to the subarray length if we fix the starting index. This property allows us to use a sliding window technique, expanding the right end of the window as long as the score remains below k, and moving the left end to reduce the sum when the score exceeds k. Since all numbers are positive, expanding the window increases the sum and the score, and shrinking the window decreases it. This guarantees that every valid subarray is counted exactly once.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n^2)	O(1)	Iterate over all subarrays, compute sum and score
Sliding Window / Two Pointers	O(n)	O(1)	Expand right pointer while score < k, shrink left pointer as needed

Algorithm Walkthrough

1. Initialize two pointers left = 0 and right = 0, a variable window_sum = 0 to keep track of the sum of the current window, and count = 0 to accumulate the number of valid subarrays.
2. Iterate right from 0 to len(nums) - 1. Add nums[right] to window_sum.
3. While the current window's score (window_sum * (right - left + 1)) >= k, shrink the window from the left by subtracting nums[left] from window_sum and incrementing left.
4. After adjusting the window, all subarrays ending at right and starting from left to right are valid. Add (right - left + 1) to count.
5. Continue this process until right reaches the end of the array.
6. Return count as the total number of subarrays satisfying the condition.

Why it works: The sliding window guarantees that the sum of the current window always produces a score strictly less than k. Since we add (right - left + 1) subarrays for each valid window ending at right, we count all subarrays exactly once, leveraging the monotonicity of the score with respect to window expansion.

Python Solution

from typing import List

class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        left = 0
        window_sum = 0
        count = 0
        for right in range(len(nums)):
            window_sum += nums[right]
            while window_sum * (right - left + 1) >= k:
                window_sum -= nums[left]
                left += 1
            count += (right - left + 1)
        return count

In the Python implementation, we maintain a running sum of the current window. The while loop ensures that the score condition is always satisfied. The (right - left + 1) term correctly counts all subarrays ending at right that satisfy the score condition.

Go Solution

func countSubarrays(nums []int, k int64) int64 {
    var left int
    var windowSum int64
    var count int64
    for right := 0; right < len(nums); right++ {
        windowSum += int64(nums[right])
        for windowSum*int64(right-left+1) >= k {
            windowSum -= int64(nums[left])
            left++
        }
        count += int64(right - left + 1)
    }
    return count
}

The Go solution mirrors the Python implementation, with explicit type conversions to int64 to handle large values and prevent integer overflow. The sliding window logic remains identical.

Worked Examples

Example 1: nums = [2,1,4,3,5], k = 10

right	window_sum	left	score = sum*len	subarrays added	count
0	2	0	2*1=2	1	1
1	3	0	3*2=6	2	3
2	7	0	7*3=21	shrink left	1
3	7+3=10	1	10*3=30	shrink left	1
4	8+5=13	3	13*2=26	shrink left	1

Final count = 6

Example 2: nums = [1,1,1], k = 5

right	window_sum	left	score	subarrays added	count
0	1	0	1*1=1	1	1
1	2	0	2*2=4	2	3
2	3	0	3*3=9	shrink left	1
2	2	1	2*2=4	2	5

Final count = 5

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n)	Each element is added and removed at most once in the sliding window, giving linear time.
Space	O(1)	Only a few integer variables are used; no extra data structures proportional to input size.

Because the algorithm uses a two-pointer approach with a monotonic sliding window, we only traverse each element once, ensuring optimal performance.

Test Cases

# Provided examples
assert Solution().countSubarrays([2,1,4,3,5], 10) == 6  # example 1
assert Solution().countSubarrays([1,1,1], 5) == 5        # example 2

# Single element array
assert Solution().countSubarrays([5], 10) == 1           # single element less than k
assert Solution().countSubarrays([10], 10) == 0          # single element equal to k

# All elements equal
assert Solution().countSubarrays([1,1,1,1], 3) == 6      # multiple valid subarrays

# Large k
assert Solution().countSubarrays([1,2,3], 100) == 6      # all subarrays valid

# Large numbers in nums
assert Solution().countSubarrays([100000,100000], 1000000000000) == 3

Test	Why
Provided examples	Validates correctness against problem statement
Single element arrays	Checks handling of smallest inputs
All elements equal	Validates counting of multiple overlapping subarrays
Large k	Ensures algorithm does not over-restrict
Large numbers in nums	Tests handling of potential integer overflow

Edge Cases

The first edge case is an array with a single element. Here, the algorithm must correctly count either one or zero subarrays depending on whether the single element's score is below k. The second edge case is when all elements are equal. This can create overlapping valid subarrays, so the counting formula (right - left + 1) ensures all are included. The third edge case is when nums contains very large values close to 10^5 and k