Count subarrays in A with sum less than k

Question

For the question below: Given an array A of N non-negative numbers and a non-negative number B,you need to find the number of subarrays in A with a sum less than B.

I have found 2 solutions: Brute force:

let count = 0;
for (let i=0; i<A.length; i++) {
    let sum = BigInt(0);
    for (let j=i; j<A.length; j++) {
        sum = sum + BigInt(A[j]);
        if (sum < B) {
            count+=1;
        }
    }
}

return count;

Optimised solution:

const countSubarrays = (A, B) => {

    let start = 0;
    let end = 0;
    let ans = 0;
    let currentSum = 0;
    let n = A.length;

    while (end < n) {
        currentSum += A[end];
        if (currentSum >= B) {
            currentSum -= A[start];
            start++;
        }

        ans += (end - start) + 1;
        end++;
    }

    return ans;
}

const A = [2, 5, 6];
const B = 10;

const result = countSubarrays(A, B);
console.log('result: ', result);

But the result for optimised one is wrong. The formula to calculate all subarrays seems incorrect: ans += (end - start) + 1; Could anyone please let me know the correct approach.

thanks

Answer 1

This can be accomplished using polynomial time:

In JS:

function count(vec, target){
  let cnt = 0;
  while(vec.length){
    let sm = 0;
    for (let i = 0; i < vec.length & target >= sm + vec[i];  i++,cnt++)
      sm += vec[i];
    vec.splice(0,1);
  }
  return cnt;
}
console.log(count([2,5,6], 10))
console.log(count([1,2,3,4,5,6,7,8],10))

C/C++:

#include <stdio.h>

int count (int * vec, int target, int m){
  int cnt = 0;
  while(m){
    int sum = 0;
    for (int i = 0; i < m & target >= sum + vec[i];  i++,cnt++)
      sum += vec[i];
    vec++;
    m--;
  }
  return cnt;
}

int main() {
    int v[3] = {2,5,6};
    int v2[8]= {1,2,3,4,5,6,7,8};
    printf("count: %d\n", count(v, 10, 3));
    printf("count: %d\n", count(v2, 10, 8));
    return 0;
}

count: 4
count: 15

Python:

def count(vec, target):
  cnt = 0
  while vec:
    sm, i = 0, 0
    while i < len(vec) and target >= (sm +vec[i]):
        sm += vec[i]
        cnt += 1
        #print(i,"\t", cnt, "\t", sm, "\t", vec[:i+1])
        i += 1
    vec = vec[1:]
  return cnt
print(count([2,5,6],10))
print(count([1,2,3,4,5,6,7,8],10))
4
15

Answer 2

The answer provided by @onyambu solves a wrong problem. The brute force solution clearly indicates that we are computing the number of subarrays, not subsequences. A subarray is continuous. For example, [2,5] is a subarray of [2,5,6] but [2,6] is not.

The problem in your optimized solution is exactly what I described in the comment. You need to keep subtracting A[start] as long as currentSum >= B before you update your answer.

The following Java code explains the correct approach with some comments.

class Solution {
    public long countSmallerSumSubarrays(int[] A, long B) {
        //if B is 0, since A[i] >= 0, no subarray is going to have a sum < 0
        if(B == 0) {
            return 0;
        }
        long ans = 0, runningSum = 0;
        int left = 0, right = 0;
        //key idea: try to compute the number of subarrays that end at A[right] with sum < B.
        //The final answer is the sum over all possible ending indices from 0 to A.length - 1.
        while(right < A.length) {
            //add the current number into running sum
            runningSum += A[right];
            //if running sum < B, it means all subarrays with left bound index in [left, right] and right bound index as right
            //are valid
            if(runningSum < B) {
                ans += right - left + 1;
            }
            //otherwise, we need to move the left index till the running sum < B
            else {
                while(runningSum >= B) {
                    runningSum -= A[left];
                    left++;
                }
                ans += right - left + 1;
            }
            //we are done computing the number of subarrays that end at A[right] with sum < B, advance right index to the next one
            right++;
        }
        return ans;
    }
    public static void main(String[] args) {
        int[] A = new int[]{2,5,6};
        Solution solution = new Solution();
        System.out.println(solution.countSmallerSumSubarrays(A, 10));
    }
}