Taro Logo

Kth Smallest Subarray Sum

Medium
Asked by:
Profile picture
20 views
Topics:
ArraysBinary SearchTwo Pointers

Given an integer array nums of length n and an integer k, return the kth smallest subarray sum.

A subarray is defined as a contiguous non-empty sequence of elements within an array.

Example 1:

Input: nums = [2, 1, 3], k = 4
Output: 3
Explanation: The subarrays are [2], [1], [3], [2, 1], [1, 3], [2, 1, 3].
The sums are 2, 1, 3, 3, 4, 6.
The 4th smallest is 3.

Example 2:

Input: nums = [3, 3, 5, 5], k = 7
Output: 10
Explanation: The subarrays are [3], [3], [5], [5], [3, 3], [3, 5], [5, 5], [3, 3, 5], [3, 5, 5], [3, 3, 5, 5].
The sums are 3, 3, 5, 5, 6, 8, 10, 11, 13, 16.
The 7th smallest is 10.

Constraints:

  • n == nums.length
  • 1 <= n <= 105
  • 1 <= nums[i] <= 100
  • 1 <= k <= n * (n + 1) / 2

Solution


Clarifying Questions

When you get asked this question in a real-life environment, it will often be ambiguous (especially at FAANG). Make sure to ask these questions in that case:

  1. What are the constraints on the size of the input array and the value of k?
  2. Can the input array contain negative numbers, zeros, or floating-point numbers?
  3. Is the array of subarray sums sorted in ascending or descending order when finding the kth smallest?
  4. If k is larger than the number of possible subarray sums, what should I return?
  5. Are we looking for contiguous subarrays only, or can the elements be non-contiguous?

Brute Force Solution

Approach

The brute force approach to finding the Kth smallest subarray sum involves considering every possible group of consecutive numbers within the given list and calculating the sum for each group. Then, we organize those sums from smallest to largest. Finally, we pick the Kth one in the sorted list.

Here's how the algorithm would work step-by-step:

  1. First, calculate the sum of every possible group of numbers. Start with groups of one number, then groups of two consecutive numbers, then three, and so on until you've considered all possible groups.
  2. Next, collect all of those sums into one big list.
  3. Then, arrange the list of sums in order from smallest to largest.
  4. Finally, find the sum located at the Kth position in that sorted list. That's your answer.

Code Implementation

def kth_smallest_subarray_sum_brute_force(numbers, k_value):
    subarray_sums = []

    # Iterate through all possible start indices
    for start_index in range(len(numbers)):

        current_subarray_sum = 0

        # Iterate through all possible end indices for each start index
        for end_index in range(start_index, len(numbers)):
            current_subarray_sum += numbers[end_index]
            subarray_sums.append(current_subarray_sum)

    # Sorting ensures we can find the kth smallest element directly
    subarray_sums.sort()

    # K_value is 1-based, so adjust for 0-based indexing
    return subarray_sums[k_value - 1]

Big(O) Analysis

Time Complexity
O(n² log n)The brute force approach first computes all possible subarray sums. There are approximately n*(n+1)/2 such subarrays, which is O(n²). Then, it sorts these sums, which takes O(n² log n²) time. Since log(n²) = 2 log(n), this is equivalent to O(n² log n). Finally, accessing the Kth element takes O(1) time. The dominant factor is sorting, so the overall time complexity is O(n² log n).
Space Complexity
O(N^2)The brute force approach described generates all possible subarray sums. In the worst case, there are N(N+1)/2 such subarrays, where N is the length of the input array. Each subarray sum is calculated and stored in a list. This list holds up to N(N+1)/2 integer values, leading to an auxiliary space usage proportional to N squared. Finally, the space used to sort the array will take O(N^2) space in the worst case, resulting in a space complexity of O(N^2).

Optimal Solution

Approach

The most efficient way to find the Kth smallest subarray sum is to avoid calculating every possible sum. We use a process of elimination by making educated guesses about the answer and then refining our guess until we find the correct one.

Here's how the algorithm would work step-by-step:

  1. First, figure out the smallest and largest possible subarray sums.
  2. Make a guess for what the Kth smallest sum might be – something in between the smallest and largest sums you calculated.
  3. Count how many subarrays have a sum less than or equal to your guess. This is crucial.
  4. If the count is less than K, it means your guess was too low. If the count is greater than or equal to K, it means your guess was too high.
  5. Adjust your guess based on whether it was too low or too high. You can use a method similar to searching a dictionary to quickly narrow down the possibilities.
  6. Keep guessing and adjusting until you find the value where exactly K-1 sums are smaller, and the current sum is the Kth smallest. That's your answer.

Code Implementation

def kth_smallest_subarray_sum(number_array, k_value):
    array_length = len(number_array)

    minimum_sum = min(number_array)
    maximum_sum = sum(number_array)

    while minimum_sum < maximum_sum:
        middle_value = (minimum_sum + maximum_sum) // 2
        
        # Count subarrays with sum less than or equal to middle_value.
        count = 0
        left_index = 0
        current_sum = 0

        for right_index in range(array_length):
            current_sum += number_array[right_index]

            while current_sum > middle_value:
                current_sum -= number_array[left_index]
                left_index += 1

            count += right_index - left_index + 1

        # Adjust search range based on the count.
        if count < k_value:
            # The middle_value is too small, so we search in the upper half.
            minimum_sum = middle_value + 1

        else:
            # The middle_value is too large, so we search in the lower half.
            maximum_sum = middle_value

    # The binary search converges to the Kth smallest sum.
    return minimum_sum

Big(O) Analysis

Time Complexity
O(n log(sum_range))The algorithm uses binary search to find the Kth smallest subarray sum. The binary search space is defined by sum_range, the difference between the maximum and minimum possible subarray sums. Inside the binary search, we iterate through the array to count the number of subarrays with a sum less than or equal to the current mid value. Counting the subarrays takes O(n) time. Since the binary search iterates log(sum_range) times, the overall time complexity is O(n log(sum_range)).
Space Complexity
O(1)The provided solution calculates subarray sums and counts them without storing all subarray sums explicitly. It only stores a few variables like the smallest possible sum, largest possible sum, the current guess, and the count of subarrays less than or equal to the guess. Since the number of variables used does not depend on the size of the input array (N), the auxiliary space complexity is constant.

Edge Cases

Null or empty input array
How to Handle:
Return null or empty list immediately to avoid NullPointerException or invalid operations.
Array with a single element
How to Handle:
Return the single element in a list, or return 0/null depending on problem specifics and constraints.
k is less than or equal to 0, or k is greater than the number of possible subarray sums
How to Handle:
Throw an IllegalArgumentException or return null/empty list to signal an invalid input value for k.
Array contains only identical values
How to Handle:
The number of subarray sums will be limited, so we need to efficiently compute and handle cases where k is within this range.
Array contains negative numbers
How to Handle:
Subarray sums can be negative, potentially impacting binary search logic or requiring adjustments to sum calculation to accommodate negative values.
Array contains very large positive numbers causing integer overflow when summing subarrays
How to Handle:
Use long data type to store subarray sums to prevent overflow, or use a different approach that avoids calculating large sums directly.
Maximum-sized input array (e.g., 10^5 elements), could cause time limit exceeded
How to Handle:
Ensure the algorithm scales efficiently, ideally O(n log n) or better, possibly using binary search and prefix sums to optimize runtime.
k is very large, approaching the maximum number of possible subarray sums (n*(n+1)/2)
How to Handle:
Efficiently compute just the necessary subarray sums using pruning or other optimization techniques when 'k' is extremely high.