Problem Description

You are given an integer array nums of length n and an integer k. Your task is to find the k-th smallest subarray sum.

A subarray is a contiguous sequence of elements from the array that contains at least one element. The subarray sum is the sum of all elements in that subarray.

For example, if nums = [1, 2, 3], the possible subarrays are:

• [1] with sum 1
• [2] with sum 2
• [3] with sum 3
• [1, 2] with sum 3
• [2, 3] with sum 5
• [1, 2, 3] with sum 6

If we sort all possible subarray sums in ascending order, we get: 1, 2, 3, 3, 5, 6. The k-th smallest value would be the element at position k in this sorted list (1-indexed).

The challenge is to find this k-th smallest sum efficiently without explicitly calculating and sorting all possible subarray sums, as there are n*(n+1)/2 possible subarrays which could be very large.

Intuition

The key insight is recognizing a monotonic property: if we pick a target sum value s, we can count how many subarrays have a sum less than or equal to s. As s increases, this count also increases (or stays the same). This is because all subarrays that had sum ≤ s will still have sum ≤ s+1, and potentially more subarrays will qualify.

This monotonic relationship suggests we can use binary search. Instead of finding all subarray sums and sorting them, we can binary search on the possible sum values. The search space ranges from the minimum element in the array (smallest possible subarray sum) to the sum of all elements (largest possible subarray sum).

For each candidate sum mid in our binary search, we need to determine: "How many subarrays have sum ≤ mid?" If this count is at least k, then the k-th smallest sum is at most mid, so we search in the lower half. If the count is less than k, then the k-th smallest sum must be larger than mid, so we search in the upper half.

The challenge now becomes efficiently counting subarrays with sum ≤ s. Since all elements are positive, we can use a sliding window approach. We maintain a window [j, i] and expand it by moving i to the right. When the window sum exceeds s, we shrink from the left by moving j. At each position i, all subarrays ending at i with starting positions from j to i have sum ≤ s, giving us i - j + 1 valid subarrays.

This transforms an initially complex problem of finding the k-th smallest among potentially O(n²) values into a binary search problem where each check takes O(n) time using the sliding window technique.

Learn more about Binary Search and Sliding Window patterns.

Solution Approach

The solution implements binary search combined with a sliding window technique. Let's walk through the implementation step by step:

Binary Search Setup:

We initialize the search boundaries:

• l = min(nums): The smallest possible subarray sum (a single element)
• r = sum(nums): The largest possible subarray sum (all elements)

The binary search looks for the smallest value s where the number of subarrays with sum ≤ s is at least k. We use bisect_left with a custom key function f that returns True when the count is at least k.

Helper Function f(s):

This function counts how many subarrays have sum ≤ s using a sliding window approach:

1. Initialize variables:

   • t = 0: Current window sum
   • j = 0: Left boundary of the window
   • cnt = 0: Count of valid subarrays

2. Iterate through the array with right pointer i:

   • Add nums[i] to the window sum: t += x
   • While the window sum exceeds s, shrink from the left:

     while t > s:
         t -= nums[j]
         j += 1

   • After adjustment, all subarrays ending at i with starting positions from j to i are valid
   • Add these subarrays to the count: cnt += i - j + 1
     • This counts: [j:i+1], [j+1:i+1], ..., [i:i+1]

3. Return cnt >= k to indicate if we've found at least k subarrays

Binary Search Execution:

The expression bisect_left(range(l, r + 1), True, key=f) searches for the leftmost position where f(s) returns True. This finds the smallest sum s where at least k subarrays have sum ≤ s.

The final answer is l + bisect_left(...) because bisect_left returns an index into the range, so we need to add the starting value l to get the actual sum value.

Time Complexity: O(n × log S) where n is the array length and S is the sum of all elements. The binary search performs O(log S) iterations, and each iteration calls f(s) which takes O(n) time.

Space Complexity: O(1) as we only use a constant amount of extra space for variables.

Example Walkthrough

Let's walk through finding the 4th smallest subarray sum for nums = [2, 1, 3] and k = 4.

Step 1: Initialize Binary Search Range

• l = min(nums) = 1 (smallest possible subarray sum)
• r = sum(nums) = 6 (largest possible subarray sum)
• We'll binary search for values between 1 and 6

Step 2: Binary Search Process

First iteration: mid = (1 + 6) / 2 = 3

• Count subarrays with sum ≤ 3 using sliding window:
  • Window [0,0]: sum = 2, valid subarrays: [2], count = 1
  • Window [0,1]: sum = 3, valid subarrays: [2,1] and [1], count = 1 + 2 = 3
  • Window [2,2]: sum = 3, valid subarrays: [3], count = 3 + 1 = 4
• Total count = 4 ≥ k = 4, so we have at least 4 subarrays with sum ≤ 3
• Search lower half: r = 3

Second iteration: mid = (1 + 3) / 2 = 2

• Count subarrays with sum ≤ 2 using sliding window:
  • Window [0,0]: sum = 2, valid subarrays: [2], count = 1
  • Window [1,1]: sum = 1, valid subarrays: [1], count = 1 + 1 = 2
  • Window cannot include index 2 (sum would be > 2)
• Total count = 2 < k = 4, so we need a larger sum
• Search upper half: l = 3

Third iteration: l = 3, r = 3

• Binary search converges to 3

Step 3: Verify Result All subarray sums sorted: [1, 2, 3, 3, 4, 6]

• [1]: sum = 1
• [2]: sum = 2
• [3]: sum = 3
• [2,1]: sum = 3 ← 4th smallest
• [1,3]: sum = 4
• [2,1,3]: sum = 6

The 4th smallest subarray sum is indeed 3.

The sliding window technique efficiently counts valid subarrays by maintaining a window that expands right and contracts left as needed, avoiding the need to explicitly generate all O(n²) subarrays.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n * log(sum(nums) - min(nums)))

The algorithm uses binary search on the range [min(nums), sum(nums)]. The binary search range has a size of sum(nums) - min(nums), leading to O(log(sum(nums) - min(nums))) iterations.

For each binary search iteration, the helper function f(s) is called, which uses a sliding window approach to count the number of subarrays with sum ≤ s. This sliding window traversal takes O(n) time since each element is visited at most twice (once by pointer i and once by pointer j).

Therefore, the total time complexity is O(n * log(sum(nums) - min(nums))).

Space Complexity: O(1)

The algorithm uses only a constant amount of extra space:

• Variables l, r for binary search bounds
• Variables t, j, cnt, i, x in the helper function f
• The range object in bisect_left is a generator that doesn't create the full list

No additional data structures that scale with input size are used, resulting in O(1) space complexity.

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

Common Pitfalls

1. Incorrect Handling of Negative Numbers

The most critical pitfall in this solution is that it assumes all numbers are positive. The sliding window technique used in count_subarrays_with_sum_at_most only works correctly when all elements are non-negative.

Why it fails with negative numbers:

• With negative numbers, adding more elements to a subarray can actually decrease the sum
• The condition while current_sum > target_sum may not correctly identify when to shrink the window
• We might miss valid subarrays because the window shrinking logic is based on the assumption that removing elements from the left always decreases the sum

Example of failure:

nums = [2, -1, 3], k = 2
# The sliding window approach would incorrectly count subarrays
# because when current_sum > target_sum, removing the leftmost element
# might actually increase the sum if that element is negative

Solution: For arrays containing negative numbers, use a different approach:

def kthSmallestSubarraySum(self, nums: List[int], k: int) -> int:
    n = len(nums)
    # Generate all possible subarray sums
    sums = []
  
    for i in range(n):
        current_sum = 0
        for j in range(i, n):
            current_sum += nums[j]
            sums.append(current_sum)
  
    # Sort and return k-th smallest
    sums.sort()
    return sums[k - 1]

Or use a more efficient heap-based approach with prefix sums for better time complexity.

2. Off-by-One Error in Binary Search Range

Another common mistake is incorrectly setting the binary search boundaries:

• Using min(0, min(nums)) instead of min(nums) when negative numbers are present
• This could cause the binary search to miss the correct answer if all subarray sums are positive but some elements are negative

Solution: When negative numbers are possible, adjust the boundaries:

# Consider that the minimum subarray sum might be negative
min_sum = min(min(nums), min(0, sum(nums)))
# Maximum could also be negative if all elements are negative
max_sum = max(max(nums), sum(nums))

3. Integer Overflow in Large Arrays

For very large arrays with large values, the sum might exceed integer limits in some languages (though Python handles arbitrary precision integers).

Solution: Be mindful of constraints and use appropriate data types if implementing in languages with fixed integer sizes.