Count Subarrays with Score Less Than K [Hard] — Shrinkable Window

Problem Statement

Score of a subarray = sum × length. Return count of subarrays with score < k.

Input: nums=[2,1,4,3,5], k=10 → Output: 6

Intuition

Shrinkable window: expand right, add to sum. While sum * (right-left+1) >= k, shrink from left. Valid subarrays ending at right = right - left + 1.

Solutions

Python

def countSubarrays(nums: list[int], k: int) -> int:
    left = s = ans = 0
    for right, val in enumerate(nums):
        s += val
        while s * (right - left + 1) >= k:
            s -= nums[left]; left += 1
        ans += right - left + 1
    return ans

C++

long long countSubarrays(vector<int>& nums, long long k) {
    long long left=0, s=0, ans=0;
    for(int r=0;r<nums.size();r++){
        s+=nums[r];
        while(s*(r-left+1)>=k) s-=nums[left++];
        ans+=r-left+1;
    }
    return ans;
}

Complexity

Time: O(n), Space: O(1)