Problem of The Day: Count Subarrays With Score Less Than K

Problem Statement

leetcode problem link

Intuition

At first glance, we are asked to count the number of subarrays where the product of the length of the subarray and the sum of its elements is less than k.
A brute-force solution would be to try all subarrays and check the condition, but that would be too slow.
Since the condition involves sum of elements, a sliding window approach came to mind: we can grow the window and shrink it dynamically depending on whether the condition is satisfied.

Approach

We use two pointers (start and end) to represent the current window.

• We expand the window by moving end and adding nums[end] to curr_total (the sum of elements inside the window).
• If at any point the condition (window size) * (curr_total) >= k is violated, we shrink the window by moving start forward and subtracting nums[start] from curr_total.
• After adjusting the window to be valid, the number of valid subarrays ending at index end is (end - start + 1).
• We accumulate this count into res for every end.

This way, we efficiently count all valid subarrays without checking each one individually.

Complexity

• Time complexity:
  \(O(n)\)
  Each element is added and removed from the sliding window at most once, leading to linear time complexity.

• Space complexity:
  \(O(1)\)
  Only a few variables are used regardless of input size.

Code

class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        start = 0
        N = len(nums)
        curr_total = 0
        res = 0
        for end in range(N):
            curr_total += nums[end]
            while start <= end and (end - start + 1) * curr_total >= k:
                curr_total -= nums[start]
                start += 1
            res += (end - start + 1)
        return res