{x}
blog image

Count Subarrays With Score Less Than K

The score of an array is defined as the product of its sum and its length.

  • For example, 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 subarrays of nums whose score is strictly less than k.

A subarray is a contiguous sequence of elements within an array.

 

Example 1:

Input: nums = [2,1,4,3,5], k = 10
Output: 6
Explanation:
The 6 subarrays having scores less than 10 are:
- [2] with score 2 * 1 = 2.
- [1] with score 1 * 1 = 1.
- [4] with score 4 * 1 = 4.
- [3] with score 3 * 1 = 3. 
- [5] with score 5 * 1 = 5.
- [2,1] with score (2 + 1) * 2 = 6.
Note that subarrays such as [1,4] and [4,3,5] are not considered because their scores are 10 and 36 respectively, while we need scores strictly less than 10.

Example 2:

Input: nums = [1,1,1], k = 5
Output: 5
Explanation:
Every subarray except [1,1,1] has a score less than 5.
[1,1,1] has a score (1 + 1 + 1) * 3 = 9, which is greater than 5.
Thus, there are 5 subarrays having scores less than 5.

 

Constraints:

  • 1 <= nums.length <= 105
  • 1 <= nums[i] <= 105
  • 1 <= k <= 1015

Solution Explanation for Count Subarrays With Score Less Than K

This problem asks to count the number of subarrays within a given array nums whose score is strictly less than a given integer k. The score of a subarray is calculated as the product of its sum and its length.

Two approaches are presented here: Prefix Sum + Binary Search and Two Pointers.

This approach leverages prefix sums to efficiently calculate the sum of any subarray. A prefix sum array s is created, where s[i] represents the sum of elements from nums[0] to nums[i-1].

The algorithm iterates through nums, considering each element as the potential right end of a subarray. For each element at index i, it performs a binary search to find the maximum length l such that the subarray from i-l to i has a score less than k. This condition is expressed as (s[i] - s[i-l]) * l < k. The binary search efficiently determines this length without explicitly checking all possible subarray lengths. The length l found is added to the final count.

Time Complexity: O(n log n), where n is the length of nums. The outer loop iterates n times, and the binary search within each iteration takes O(log n) time.

Space Complexity: O(n) due to the prefix sum array s.

Approach 2: Two Pointers (Sliding Window)

This approach uses a sliding window technique with two pointers, i and j. i represents the right boundary of the window, and j represents the left boundary. The window continuously expands to the right (increasing i). As the window grows, the score is checked. If the score becomes greater than or equal to k, the window's left boundary is shifted right (increasing j) until the score condition is met again. The number of subarrays ending at i is simply the window size (i - j + 1).

Time Complexity: O(n). The pointers i and j each traverse the array at most once.

Space Complexity: O(1). Only a constant number of variables are used.

Code Examples

The code examples below demonstrate both approaches in Python, Java, C++, and Go. The Two Pointers approach is generally preferred due to its linear time complexity.

Python3 (Prefix Sum + Binary Search)

from itertools import accumulate
 
class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        s = list(accumulate(nums, initial=0))
        ans = 0
        for i in range(1, len(s)):
            left, right = 0, i
            while left < right:
                mid = (left + right + 1) // 2
                if (s[i] - s[i - mid]) * mid < k:
                    left = mid
                else:
                    right = mid - 1
            ans += left
        return ans

Python3 (Two Pointers)

class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        ans = s = j = 0
        for i, v in enumerate(nums):
            s += v
            while s * (i - j + 1) >= k:
                s -= nums[j]
                j += 1
            ans += i - j + 1
        return ans

Java (Two Pointers)

class Solution {
    public long countSubarrays(int[] nums, long k) {
        long ans = 0, s = 0;
        for (int i = 0, j = 0; i < nums.length; ++i) {
            s += nums[i];
            while (s * (i - j + 1) >= k) {
                s -= nums[j++];
            }
            ans += i - j + 1;
        }
        return ans;
    }
}

C++ (Two Pointers)

class Solution {
public:
    long long countSubarrays(vector<int>& nums, long long k) {
        long long ans = 0, s = 0;
        for (int i = 0, j = 0; i < nums.size(); ++i) {
            s += nums[i];
            while (s * (i - j + 1) >= k) {
                s -= nums[j++];
            }
            ans += i - j + 1;
        }
        return ans;
    }
};

Go (Two Pointers)

func countSubarrays(nums []int, k int64) (ans int64) {
	s, j := 0, 0
	for i, v := range nums {
		s += v
		for int64(s*(i-j+1)) >= k {
			s -= nums[j]
			j++
		}
		ans += int64(i - j + 1)
	}
	return
}

The Java, C++, and Go examples for the Two Pointers approach are very similar in structure to the Python example. They all efficiently solve the problem in linear time. Remember to choose the approach that best suits your needs and coding style, keeping in mind the time and space complexity trade-offs.