Kth Smallest sum of continuous subarrays of positive numbers

Given an sorted array of positive numbers our tasks is to find the kth smallest sum of continuous subarray. Examples:

Input : a[] = {1, 2, 3, 4, 5, 6} k = 4 
Output : 3 Explanation : List of sorted subarray sum: {1, 2, 3, 3, 4, 5, 5, 6, 6, 7, 9, 9, 10, 11, 12, 14, 15, 15, 18, 20, 21}. 4th smallest element is 3. 

Input : a[] = {1, 2, 3, 4, 5, 6} k = 13 
Output: 10

A naive approach is to first generate all the continuous subarray sums which can be done in O(N^2) by precomputing prefix sum. Sort the sum array and give the kth smallest element. 

A better approach (for arrays with small sum) is to use Binary search. First, we will precompute a prefix sum array. Now we apply binary search on the number which are possible candidates for the kth smallest sum which will be in range [0, total sum of the array]. For now let’s assume we have a function called calculateRank which will give us the rank of any number in the sorted array of continuous subarray sum. 

In binary search we will use this calculateRank function to check if the rank of the mid element is less than K, if yes then we reduce the start point to mid+1 else if it is greater than or equal to K, then reduce the end point to mid-1 and also updating the answer variable. Now let’s come back to calculateRank function. Here, also we will use binary search but on our prefix sum array.We will iterate on our array and lets assume we are on the ith index, we will count how many subarrays can be made with starting element as this ith element whose sum is less than the mid element whose rank we have to calculate. 

We do for all elements and add the count of each which we the rank of the mid-element. To count the number of the subarray starting with the ith index we use apply binary on prefix sum. C++ implementation to make things clearer. 

Implementation:

C++

// CPP program to find k-th smallest sum
#include "algorithm"
#include "iostream"
#include "vector"
using namespace std;

int CalculateRank(vector<int> prefix, int n, int x)
{
    int cnt;

    // Initially rank is 0.
    int rank = 0;
    int sumBeforeIthindex = 0;
    for (int i = 0; i < n; ++i) {

        // Calculating the count the subarray with
        // starting at ith index
        cnt = upper_bound(prefix.begin(), prefix.end(),
                sumBeforeIthindex + x) - prefix.begin();

        // Subtracting the subarrays before ith index.
        cnt -= i;

        // Adding the count to rank.
        rank += cnt;
        sumBeforeIthindex = prefix[i];
    }
    return rank;
}

int findKthSmallestSum(int a[], int n, int k)
{
    // PrefixSum array.
    vector<int> prefix;

    // Total Sum initially 0.
    int sum = 0;
    for (int i = 0; i < n; ++i) {
        sum += a[i];
        prefix.push_back(sum);
    }

    // Binary search on possible
    // range i.e [0, total sum]
    int ans = 0;
    int start = 0, end = sum;
    while (start <= end) {

        int mid = (start + end) >> 1;

        // Calculating rank of the mid and
        // comparing with K
        if (CalculateRank(prefix, n, mid) >= k) {

            // If greater or equal store the answer
            ans = mid;
            end = mid - 1;
        }
        else {
            start = mid + 1;
        }
    }

    return ans;
}

int main()
{
    int a[] = { 1, 2, 3, 4, 5, 6 };
    int k = 13;
    int n = sizeof(a)/sizeof(a[0]);
    cout << findKthSmallestSum(a, n, k);
    return 0;
}

Java

// Java program to find k-th smallest sum
import java.util.*;

public class Main {
    public static int
    calculateRank(ArrayList<Integer> prefix, int n, int x)
    {
        int cnt;

        // Initially rank is 0.
        int rank = 0;
        int sumBeforeIthindex = 0;
        for (int i = 0; i < n; ++i) {

            // Calculating the count the subarray with
            // starting at ith index
            cnt = Collections.binarySearch(
                prefix, sumBeforeIthindex + x);

            if (cnt >= 0) {
                // If the searched element is found,
                // move to the last occurrence of the
                // element
                while (cnt < n
                       && prefix.get(cnt)
                              == sumBeforeIthindex + x) {
                    cnt++;
                }
            }
            else {
                // If the searched element is not found,
                // take the bitwise complement of the index
                // to get the position of the first element
                // greater than the searched element
                cnt = ~cnt;
            }

            // Subtracting the subarrays before ith index.
            cnt -= i;

            // Adding the count to rank.
            rank += cnt;
            sumBeforeIthindex = prefix.get(i);
        }
        return rank;
    }

    public static int findKthSmallestSum(int a[], int n,
                                         int k)
    {
        // PrefixSum array.
        ArrayList<Integer> prefix = new ArrayList<>();

        // Total Sum initially 0.
        int sum = 0;
        for (int i = 0; i < n; ++i) {
            sum += a[i];
            prefix.add(sum);
        }

        // Binary search on possible
        // range i.e [0, total sum]
        int ans = 0;
        int start = 0, end = sum;
        while (start <= end) {

            int mid = (start + end) >> 1;

            // Calculating rank of the mid and
            // comparing with K
            if (calculateRank(prefix, n, mid) >= k) {

                // If greater or equal store the answer
                ans = mid;
                end = mid - 1;
            }
            else {
                start = mid + 1;
            }
        }

        return ans;
    }

    public static void main(String[] args)
    {
        int a[] = { 1, 2, 3, 4, 5, 6 };
        int k = 13;
        int n = a.length;
        System.out.println(findKthSmallestSum(a, n, k));
    }
}

Python3

# Python3 program to find k-th smallest sum

from bisect import bisect_right

def CalculateRank(prefix, n, x):
    # Initially rank is 0.
    rank = 0
    sumBeforeIthindex = 0
    for i in range(n):
        # Calculating the count of the subarray with starting at ith index
        cnt = bisect_right(prefix, sumBeforeIthindex + x)

        # Subtracting the subarrays before ith index.
        cnt -= i

        # Adding the count to rank.
        rank += cnt
        sumBeforeIthindex = prefix[i]
    return rank

def findKthSmallestSum(a, n, k):
    # PrefixSum array.
    prefix = []

    # Total Sum initially 0.
    sum = 0
    for i in range(n):
        sum += a[i]
        prefix.append(sum)

    # Binary search on possible range i.e [0, total sum]
    ans = 0
    start = 0
    end = sum
    while start <= end:
        mid = (start + end) // 2

        # Calculating rank of the mid and comparing with K
        if CalculateRank(prefix, n, mid) >= k:
            # If greater or equal store the answer
            ans = mid
            end = mid - 1
        else:
            start = mid + 1

    return ans

# Example usage
a = [1, 2, 3, 4, 5, 6]
k = 13
n = len(a)
print(findKthSmallestSum(a, n, k)) # Output: 24

Javascript

function CalculateRank(prefix, n, x)
{

    // Initially rank is 0.
    let rank = 0;
    let sumBeforeIthindex = 0;
    for (let i = 0; i < n; i++)
    {

        // Calculating the count the subarray with
        // starting at ith index
        let cnt = prefix.filter(val => val <= sumBeforeIthindex + x).length - i;

        // Adding the count to rank.
        rank += cnt;
        sumBeforeIthindex = prefix[i];
    }
    return rank;
}

function findKthSmallestSum(a, n, k) {

// PrefixSum array
    let prefix = [];

    // Total Sum initially 0.
    let sum = 0;
    for (let i = 0; i < n; i++) {

    // Binary search on possible
    // range i.e [0, total sum]
        sum += a[i];
        prefix.push(sum);
    }

    let ans = 0;
    let start = 0;
    let end = sum;
    while (start <= end) {
        let mid = Math.floor((start + end) / 2);

        // Calculating rank of the mid and
        // comparing with K
        if (CalculateRank(prefix, n, mid) >= k) {

        // If greater or equal store the answer
            ans = mid;
            end = mid - 1;
        } else {
            start = mid + 1;
        }
    }

    return ans;
}

// Example usage
let a = [1, 2, 3, 4, 5, 6]
let k = 13
let n = a.length
console.log(findKthSmallestSum(a, n, k))

C#

using System;
using System.Linq;

namespace ConsoleApp
{
    class Program
    {
        static int CalculateRank(int[] prefix, int n, int x)
        {
            // Initially rank is 0.
            int rank = 0;
            int sumBeforeIthindex = 0;
            for (int i = 0; i < n; i++)
            {
                // Calculating the count the subarray with
                // starting at ith index
                int cnt = prefix.Count(val => val <= sumBeforeIthindex + x) - i;

                // Adding the count to rank.
                rank += cnt;
                sumBeforeIthindex = prefix[i];
            }
            return rank;
        }

        static int FindKthSmallestSum(int[] a, int n, int k)
        {
            // PrefixSum array
            int[] prefix = new int[n];

            // Total Sum initially 0.
            int sum = 0;
            for (int i = 0; i < n; i++)
            {
                // Binary search on possible
                // range i.e [0, total sum]
                sum += a[i];
                prefix[i] = sum;
            }

            int ans = 0;
            int start = 0;
            int end = sum;
            while (start <= end)
            {
                int mid = (start + end) / 2;

                // Calculating rank of the mid and
                // comparing with K
                if (CalculateRank(prefix, n, mid) >= k)
                {
                    // If greater or equal store the answer
                    ans = mid;
                    end = mid - 1;
                }
                else
                {
                    start = mid + 1;
                }
            }

            return ans;
        }

        static void Main(string[] args)
        {
            // Example usage
            int[] a = {1, 2, 3, 4, 5, 6};
            int k = 13;
            int n = a.Length;
            Console.WriteLine(FindKthSmallestSum(a, n, k));
        }
    }
}

Output

10

Time Complexity: O(N Log N Log SUM). Here N is number of elements and SUM is array sum. CalculateRank function takes O(N log N) time and is called Log SUM times.