Problem Description

You are given a positive integer array nums and an integer k. The score of an array is calculated by multiplying the sum of its elements by its length.

For example, if we have an array [1, 2, 3, 4, 5]:

• Sum of elements = 1 + 2 + 3 + 4 + 5 = 15
• Length = 5
• Score = 15 × 5 = 75

Your task is to count how many non-empty subarrays of nums have a score that is strictly less than k.

A subarray is a contiguous sequence of elements from the original array. For instance, if nums = [1, 2, 3, 4], then [2, 3] and [1, 2, 3] are subarrays, but [1, 3] is not (elements are not contiguous).

The function should return the total count of all such subarrays whose score is less than k.

Intuition

When we need to count subarrays with a certain property, a brute force approach would be to check every possible subarray. However, this would be inefficient as there are O(n²) subarrays.

Let's think about the structure of the problem. For a subarray from index i to j, the score is (sum of elements) × (length). If we fix the ending position j, we need to find all valid starting positions i where the score is less than k.

Notice an important property: as we extend a subarray to the left (decreasing the starting index while keeping the ending fixed), both the sum and the length increase. This means the score increases monotonically. In other words, if a subarray of length L ending at position j has a score ≥ k, then all longer subarrays ending at j will also have scores ≥ k.

This monotonic property is perfect for binary search! For each ending position, we can use binary search to find the maximum length l such that the subarray of length l ending at that position has a score less than k. Then all subarrays of length 1, 2, ..., l ending at that position are valid.

To efficiently calculate the sum of any subarray, we can use prefix sums. With prefix sum array s, the sum of elements from index i to j is simply s[j+1] - s[i].

So the approach becomes: for each ending position, binary search for the maximum valid length, and add that length to our answer. The sum of all these lengths gives us the total count of valid subarrays.

Learn more about Binary Search, Prefix Sum and Sliding Window patterns.

Solution Approach

The solution uses prefix sum combined with binary search to efficiently count the valid subarrays.

Step 1: Build Prefix Sum Array

First, we create a prefix sum array s where s[i] represents the sum of the first i elements of nums. We initialize it with 0 at index 0 for easier calculation. This allows us to compute the sum of any subarray in O(1) time: the sum from index i to j is s[j+1] - s[i].

s = list(accumulate(nums, initial=0))

Step 2: Enumerate Ending Positions

We iterate through each position i (from 1 to len(s)-1) as the ending position of potential subarrays. Position i in the prefix sum array corresponds to including the first i elements of the original array.

Step 3: Binary Search for Maximum Valid Length

For each ending position i, we need to find the maximum length l such that a subarray of length l ending at position i-1 in the original array has a score less than k.

The binary search works as follows:

• Search range: l from 0 to i (the maximum possible length)
• For a candidate length mid:
  • The sum of the subarray is s[i] - s[i - mid]
  • The score is (s[i] - s[i - mid]) * mid
  • If score < k, we can try a longer subarray (set l = mid)
  • Otherwise, we need a shorter subarray (set r = mid - 1)

while l < r:
    mid = (l + r + 1) >> 1  # equivalent to (l + r + 1) // 2
    if (s[i] - s[i - mid]) * mid < k:
        l = mid
    else:
        r = mid - 1

Step 4: Count Valid Subarrays

After the binary search, l represents the maximum length where the score is less than k. This means all subarrays of length 1, 2, ..., l ending at this position are valid, contributing l to our answer.

Step 5: Return Total Count

We sum up the counts for all ending positions to get the final answer.

The time complexity is O(n log n) where n is the length of the array, as we perform a binary search for each of the n positions. The space complexity is O(n) for storing the prefix sum array.

Example Walkthrough

Let's walk through the solution with nums = [2, 1, 4, 3, 5] and k = 10.

Step 1: Build Prefix Sum Array

• Original array: [2, 1, 4, 3, 5]
• Prefix sum s: [0, 2, 3, 7, 10, 15]
  • s[0] = 0 (initial value)
  • s[1] = 0 + 2 = 2
  • s[2] = 2 + 1 = 3
  • s[3] = 3 + 4 = 7
  • s[4] = 7 + 3 = 10
  • s[5] = 10 + 5 = 15

Step 2-4: Process Each Ending Position

Position i = 1 (subarrays ending at index 0 of original array):

• Binary search for max length where score < 10
• Try length 1: sum = s[1] - s[0] = 2 - 0 = 2, score = 2 × 1 = 2 < 10 ✓
• Maximum valid length = 1
• Valid subarrays: [2]
• Count: 1

Position i = 2 (subarrays ending at index 1):

• Binary search range: length 0 to 2
• Try length 2: sum = s[2] - s[0] = 3 - 0 = 3, score = 3 × 2 = 6 < 10 ✓
• Maximum valid length = 2
• Valid subarrays: [1] (length 1), [2,1] (length 2)
• Count: 2

Position i = 3 (subarrays ending at index 2):

• Binary search range: length 0 to 3
• Try length 2: sum = s[3] - s[1] = 7 - 2 = 5, score = 5 × 2 = 10 ≥ 10 ✗
• Try length 1: sum = s[3] - s[2] = 7 - 3 = 4, score = 4 × 1 = 4 < 10 ✓
• Maximum valid length = 1
• Valid subarrays: [4]
• Count: 1

Position i = 4 (subarrays ending at index 3):

• Binary search range: length 0 to 4
• Try length 2: sum = s[4] - s[2] = 10 - 3 = 7, score = 7 × 2 = 14 ≥ 10 ✗
• Try length 1: sum = s[4] - s[3] = 10 - 7 = 3, score = 3 × 1 = 3 < 10 ✓
• Maximum valid length = 1
• Valid subarrays: [3]
• Count: 1

Position i = 5 (subarrays ending at index 4):

• Binary search range: length 0 to 5
• Try length 3: sum = s[5] - s[2] = 15 - 3 = 12, score = 12 × 3 = 36 ≥ 10 ✗
• Try length 1: sum = s[5] - s[4] = 15 - 10 = 5, score = 5 × 1 = 5 < 10 ✓
• Maximum valid length = 1
• Valid subarrays: [5]
• Count: 1

Step 5: Total Count Total = 1 + 2 + 1 + 1 + 1 = 6

The 6 valid subarrays are: [2], [1], [2,1], [4], [3], [5]

Solution Implementation

Time and Space Complexity

Time Complexity: O(n × log n)

The algorithm uses a prefix sum array and then iterates through each position i from 1 to n (where n is the length of the array). For each position i, it performs a binary search to find the maximum length l of a subarray ending at position i-1 that satisfies the condition (s[i] - s[i - mid]) * mid < k.

• The outer loop runs n times
• For each iteration, the binary search performs O(log i) operations, which in the worst case is O(log n)
• Therefore, the overall time complexity is O(n × log n)

Space Complexity: O(n)

The algorithm creates a prefix sum array s using accumulate with initial=0, which has length n + 1. This is the only additional data structure that scales with the input size. The other variables (ans, i, l, r, mid) use constant space.

• The prefix sum array s requires O(n) space
• All other variables use O(1) space
• Therefore, the overall space complexity is O(n)

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Binary Search Boundary Calculation Error

The Pitfall: The most common mistake is using the wrong midpoint calculation in binary search. When searching for the maximum valid length, using mid = (left + right) // 2 instead of mid = (left + right + 1) // 2 will cause an infinite loop when left = right - 1 and the condition is satisfied.

Why it happens:

• When left = right - 1 and we use mid = (left + right) // 2, mid equals left
• If the condition (s[i] - s[i - mid]) * mid < k is true, we set left = mid, which doesn't change left
• This creates an infinite loop

Solution: Always use mid = (left + right + 1) // 2 (or (left + right + 1) >> 1) when searching for the maximum value where we update left = mid on success.

2. Score Calculation Overflow

The Pitfall: For large arrays with large values, calculating subarray_sum * length might cause integer overflow in languages with fixed integer sizes. While Python handles arbitrary precision integers automatically, this is a critical issue in languages like Java or C++.

Solution:

• In Python: No action needed due to automatic handling
• In other languages: Use long/int64 types or check for overflow before multiplication
• Alternative approach: Rearrange the inequality to avoid multiplication: instead of sum * length < k, check sum < k / length (being careful with integer division)

3. Off-by-One Error in Prefix Sum Indexing

The Pitfall: Confusing the indexing between the original array and the prefix sum array. The prefix sum array has length n+1 where prefix_sums[i] represents the sum of the first i elements of the original array.

Why it happens:

• prefix_sums[0] = 0 (sum of zero elements)
• prefix_sums[i] - prefix_sums[j] gives the sum of elements from index j to i-1 in the original array
• It's easy to miscalculate which elements are included in the subarray

Solution:

• Always remember that prefix_sums[end] - prefix_sums[start] gives the sum of nums[start:end]
• The length of this subarray is end - start, not end - start + 1
• Draw out small examples to verify the indexing

4. Incorrect Initialization of Binary Search Bounds

The Pitfall: Setting left = 1 instead of left = 0 in the binary search, which would miss counting subarrays of length 0 (though in this problem, we want non-empty subarrays, so length must be at least 1).

Solution:

• Initialize left = 0 to handle edge cases properly
• The binary search will naturally find left = 0 if no valid subarrays exist
• If you need to ensure non-empty subarrays, you can either:
  • Add a check after binary search: total_count += max(0, left)
  • Or initialize left = 1 but handle the case where even length-1 subarray exceeds k