Smallest subarray with positive sum for all indices

Given an array arr[] of size N. The task is to determine the minimum length of a subarray starting from index i, such that the sum of the subarray is strictly greater than 0. Calculate the length for all i's in the range 1 to N. If no such subarray exists, then return 0.

Examples:

Input: N = 3, arr[] = {-4, 3, 4}
Output: {3, 1, 1}
Explanation: For index 0, we need to take the whole array for the sum to be positive, whereas for the other two indices taking only a single element makes the sum positive.

Input: N = 5, arr[] = {-2, -3, 5, -2, 1}
Output: {0, 2, 1, 0, 1}
Explanation: For indexes 0 and 3, no subarray with a positive sum exists.

Approach: To solve the problem, follow the below idea:

To find the minimum length of subarrays with sums greater than 0 for each index in the given array, use a prefix sum array and a stack to track the next greater element. Finding the next greater element in the prefix sum array will tell us the smallest length of subarray with positive sum.

It can be observed that for any index i, if the next greater element in the prefix sum array is at index j, then all the subarrays which start at index i and end at index < j will have negative sum.

Iterate through the array, updating the result array accordingly. The final result, excluding the last element, yields the desired lengths.

Step-by-step algorithm:

• Prefix Sum:
  • Create an array (pref) where each element represents the cumulative sum of elements up to that index.
• Initialize:
  • Set up a stack(s) to keep track of indices and initialize a result array (res).
• Find Length for First Element:
  • Iterate from left to right in the array.
  • If the prefix sum at an index is greater than 0, set the first element in the result array (res[0]) to the current index + 1.
• Next Greater Element:
  • Iterate from right to left in the array.
  • Use the stack to efficiently find the next greater element for each index.
  • Update the result array based on the stack.
• Final Result:
  • Remove the last element from the result array to get the final lengths.

Below is the implementation of the algorithm:

#include <iostream>
#include <vector>
#include <stack>

using namespace std;

// Function to find minimum length for each element
vector<int> findAnswer(int N, vector<int>& arr) {

    // Create a pref array
    vector<long long> pref(N, arr[0]);

    // Iterate in array
    for (int i = 1; i < N; i++) {
        pref[i] = pref[i - 1] + arr[i];
    }

    // Initialise a stack
    stack<int> s;
    vector<int> res(N + 1, 0);

    // Iterate in array to get the length for first element
    for (int i = 0; i < N; i++) {
        if (pref[i] > 0) {
            res[0] = i + 1;
            break;
        }
    }

    // Iterate from backward to get the prefix length for each element
    for (int i = N - 1; i >= 0; i--) {

        // If stack is not empty
        while (!s.empty() && pref[s.top()] <= pref[i])
            s.pop();

        if (s.empty())
            res[i + 1] = 0;
        else if (i < N - 1)
            res[i + 1] = s.top() - i;

        s.push(i);
    }

    // Remove the last element
    res.pop_back();
    return res;
}

// Driver code
int main() {
    vector<int> arr = { -2, -3, 5, -2, 1 };
    int N = arr.size();

    // Function call
    vector<int> ans = findAnswer(N, arr);

    for (auto a : ans) {
        cout << a << " ";
    }
    return 0;
}

import java.util.*;

public class Solution {
    // Function to find minimum length for each element
    static List<Integer> findAnswer(int N, List<Integer> arr) {
        // Create a pref array
        List<Long> pref = new ArrayList<>(Collections.nCopies(N, 0L));

        // Iterate in array
        pref.set(0, (long) arr.get(0));
        for (int i = 1; i < N; i++) {
            pref.set(i, pref.get(i - 1) + arr.get(i));
        }

        // Initialise a stack
        Stack<Integer> s = new Stack<>();
        List<Integer> res = new ArrayList<>(Collections.nCopies(N + 1, 0));

        // Iterate in array to get the length for the first element
        for (int i = 0; i < N; i++) {
            if (pref.get(i) > 0) {
                res.set(0, i + 1);
                break;            }
        }

        // Iterate from backward to get the prefix length for each element
        for (int i = N - 1; i >= 0; i--) {
            // If stack is not empty
            while (!s.isEmpty() && pref.get(s.peek()) <= pref.get(i))
                s.pop();

            if (s.isEmpty())
                res.set(i + 1, 0);
            else if (i < N - 1)
                res.set(i + 1, s.peek() - i);

            s.push(i);
        }

        // Remove the last element
        res.remove(res.size() - 1);
        return res;
    }

    // Driver code
    public static void main(String[] args) {
        List<Integer> arr = Arrays.asList(-2, -3, 5, -2, 1);
        int N = arr.size();

        // Function call
        List<Integer> ans = findAnswer(N, arr);

        for (int a : ans) {
            System.out.print(a + " ");
        }
    }
}

// This code is contributed by shivamgupta0987654321

# Python Implementation

# Function to find minimum length for each element
def findAnswer(N, arr):
    # Create a prefix sum array
    pref = [arr[0]] * N

    # Calculate prefix sums
    for i in range(1, N):
        pref[i] = pref[i - 1] + arr[i]

    # Initialize a stack and result array
    s = []
    res = [0] * (N + 1)

    # Iterate to get the length for the first element
    for i in range(N):
        if pref[i] > 0:
            res[0] = i + 1
            break

    # Iterate backward to get the prefix length for each element
    for i in range(N - 1, -1, -1):
        # Pop elements from the stack if prefix sum is greater or equal
        while s and pref[s[-1]] <= pref[i]:
            s.pop()

        if not s:
            res[i + 1] = 0
        elif i < N - 1:
            res[i + 1] = s[-1] - i

        s.append(i)

    # Remove the last element
    res.pop()
    return res

# Main function
arr = [-2, -3, 5, -2, 1]
N = len(arr)

# Call the function and print the result
ans = findAnswer(N, arr)

for a in ans:
    print(a, end=" ")
    
# This code is contributed by Tapesh(tapeshdu420)

using System;
using System.Collections.Generic;

public class Solution {
    // Function to find minimum length for each element
    static List<int> FindAnswer(int N, List<int> arr)
    {
        // Create a pref array
        List<long> pref = new List<long>(new long[N]);
        pref[0] = arr[0];

        // Iterate in array
        for (int i = 1; i < N; i++) {
            pref[i] = pref[i - 1] + arr[i];
        }

        // Initialise a stack
        Stack<int> s = new Stack<int>();
        List<int> res = new List<int>(new int[N + 1]);

        // Iterate in array to get the length for first
        // element
        for (int i = 0; i < N; i++) {
            if (pref[i] > 0) {
                res[0] = i + 1;
                break;
            }
        }

        // Iterate from backward to get the prefix length
        // for each element
        for (int i = N - 1; i >= 0; i--) {
            // If stack is not empty
            while (s.Count > 0 && pref[s.Peek()] <= pref[i])
                s.Pop();

            if (s.Count == 0)
                res[i + 1] = 0;
            else if (i < N - 1)
                res[i + 1] = s.Peek() - i;

            s.Push(i);
        }

        // Remove the last element
        res.RemoveAt(N);
        return res;
    }

    // Driver code
    public static void Main(string[] args)
    {
        List<int> arr
            = new List<int>() { -2, -3, 5, -2, 1 };
        int N = arr.Count;

        // Function call
        List<int> ans = FindAnswer(N, arr);

        foreach(var a in ans) { Console.Write(a + " "); }
    }
}

// Function to find minimum length for each element
function findAnswer(arr) {
    const N = arr.length;

    // Create a pref array
    const pref = new Array(N).fill(arr[0]);

    // Iterate in array
    for (let i = 1; i < N; i++) {
        pref[i] = pref[i - 1] + arr[i];
    }

    // Initialise a stack
    const s = [];
    const res = new Array(N + 1).fill(0);

    // Iterate in array to get the length for first element
    for (let i = 0; i < N; i++) {
        if (pref[i] > 0) {
            res[0] = i + 1;
            break;
        }
    }

    // Iterate from backward to get the prefix length for each element
    for (let i = N - 1; i >= 0; i--) {
        // If stack is not empty
        while (s.length && pref[s[s.length - 1]] <= pref[i]) {
            s.pop();
        }

        if (!s.length) {
            res[i + 1] = 0;
        } else if (i < N - 1) {
            res[i + 1] = s[s.length - 1] - i;
        }

        s.push(i);
    }

    // Remove the last element
    res.pop();
    return res;
}

// Driver code
const arr = [-2, -3, 5, -2, 1];

// Function call
const ans = findAnswer(arr);

console.log(ans.join(" "));

0 2 1 0 1 

Time Complexity: O(N), where N is the length of input array arr[].
Auxiliary Space: O(N)