Count Subarrays With Score Less Than K

hard sliding window two pointers subarray

Problem

The score of an array is its sum multiplied by its length — e.g. the score of [1,2,3,4,5] is (1+2+3+4+5)·5 = 75. Given a positive integer array nums and an integer k, return the number of non-empty contiguous subarrays whose score is strictly less than k.

Inputnums = [2,1,4,3,5], k = 10

Output6

The 6 good subarrays are [2], [1], [4], [3], [5] (each length 1) and [2,1] with score (2+1)·2 = 6. [1,4] scores 10 and [4,3,5] scores 36, both excluded.

Inputnums = [1,1,1], k = 5

Output5

Every subarray except [1,1,1] qualifies; [1,1,1] scores (1+1+1)·3 = 9 ≥ 5.

nums (comma-separated, positive)

k (integer)

def count_subarrays(nums, k):
    ans = 0
    total = 0
    left = 0
    for right in range(len(nums)):
        total += nums[right]
        while total * (right - left + 1) >= k:
            total -= nums[left]
            left += 1
        ans += right - left + 1
    return ans

function countSubarrays(nums, k) {
  let ans = 0, total = 0, left = 0;
  for (let right = 0; right < nums.length; right++) {
    total += nums[right];
    while (total * (right - left + 1) >= k) {
      total -= nums[left];
      left++;
    }
    ans += right - left + 1;
  }
  return ans;
}

long countSubarrays(int[] nums, long k) {
    long ans = 0, total = 0;
    int left = 0;
    for (int right = 0; right < nums.length; right++) {
        total += nums[right];
        while (total * (right - left + 1) >= k) {
            total -= nums[left];
            left++;
        }
        ans += right - left + 1;
    }
    return ans;
}

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

Explanation

Brute force would test every subarray, which is O(n²). The key observation that unlocks a linear sliding window is monotonicity: every element is positive, so extending a window can only raise its sum and its length, hence its score never decreases. Likewise, shrinking the window from the left can only lower the score.

So we maintain a window [left .. right] together with its running sum total. The window's score is total · (right − left + 1). We advance right one step at a time, adding nums[right] to total.

Whenever the score reaches k or more, the window is too "heavy", so we drop elements from the left — subtracting nums[left] and incrementing left — until the score is once again below k.

At that point [left .. right] is the longest valid window ending at right. Because of monotonicity, every subarray ending at right that starts at or after left is also valid. There are right − left + 1 of them, so we add that count to ans.

Each of left and right only moves forward across the whole array, giving O(n) time. Note the score can be large, so the running products use 64-bit integers in Java and C++.

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