Last Updated: 13 Nov, 2021

Equilibrium

Hard
Asked in company
Uber

Problem statement

There are ‘N’ water balls arranged in a line. Some of them are hot water balls, and some are cold, i.e., temperatures of balls may be negative, zero or positive. You have an empty bag with you, and you will be collecting some balls in this bag, starting from left to right.

For each ball, you decide whether to collect this ball or not. Once you collect a ball, you can’t drop it, and you can’t go left from any position.

Total temperature of the bag at any moment is the sum of temperatures of water balls collected in it till that moment. Your task is to find the maximum number of balls you can collect, such that the total temperature of your bag must not go negative at any instance.

Input Format:
The first line contains ‘T’, denoting the number of tests.
For each Test :
    The first line contains an integer ‘N’, denoting the number of water balls.
    The second line contains an array ‘A’ of length ‘N’, denoting the temperatures of water balls.
Output Format:
For each test, print an integer, denoting the maximum number of balls you can collect, such that the total temperature of the bag is non-negative at each instance.
Note:
You are not required to print the expected output. It has already been taken care of. Just implement the function.
Constraints -
1 <= T <= 5
1 <= N <= 10^5
-10^9 <= A[i] <= 10^9  i ∈  (1,N)


Time Limit: 1 sec

Approaches

01 Approach

Approach:  We need to find the longest subsequence such that its prefix sum at each position is non-negative. Using recursion, generate all subsequences of the given array. For each subsequence, check whether its prefix sum is non-negative at each place. If the prefix sum goes negative, ignore such a sequence; otherwise, consider its length to maximize the answer.


 

Algorithm:

  • Implement a recursive function -
    • That takes some parameters as follows -
      • Current index of the array is to be explored in subsequences.
      • Prefix sum of already existing subsequence in which current index will be used.
      • Input array and its size.
    • Returns length of possible subsequence, starting from ‘index’ (passed in argument) to end of the array.
  • Function implementation goes as follows -
    • If the prefix sum of subsequence is already negative, can't process the remaining sequence, return -1.
    • If we have reached the end of the array, then the length of the remaining sequence is 0; hence return 0.
    • Call two recursive calls to the next index as follows:
      • With adding current value to existing prefix sum. Here current value is considered; hence the answer for this call will be 1+ [next recursive call].
      • Another is without taking the current value. Here the answer will be the same from [next recursive call].
      • Return maximum answer among both cases.
  • Call the above recursion from the main solution function, for index = 0 and prefix_sum = 0. This call returns our desired answer.

02 Approach

Approach: Above discussed brute force approach can be optimized using dynamic programming.

Maintain a dp[ ][ ], where dp[i][j] contains the maximum possible prefix sum if ‘j’ elements are taken till index ‘i’. Here index can vary from [0 to N-1] and the count of collected elements from [0 to N], so dimensions of dp will be (N * (N+1)). Initialize the entire dp with some negative value, say -1.

Now the question arises, what are the transitions to fill the dp?

  1. Let’s fill dp[i][j] without including i-th element, then maximum prefix sum of [‘j’ elements taken till i-th index] will be the same as [‘j’ elements taken till (i-1)’th index].
  2. Let’s fill dp[i][j] including i-th element, then maximum prefix sum of [‘j’ elements taken till i-th index] will be the sum of [‘j-1’ elements taken till (i-1)’th index] and current value A[i].

Do the above transitions for each pair (i, j), and our desired answer will be the maximum value of ‘j’ with i = N-1, for which dp[i][j] is non-negative.

 

Algorithm:

  • Initialize a dp[N][N + 1] with -1.
  • Initialize dp[i][0] with zero for all [i = 0 to N-1]. This denotes the length of subsequence till index ‘i’ when 0 elements are taken.
  • Solve for the first element manually, i.e, if A[0] is non-negative, dp[0][1] = A[0].
  • Iterate i = 1 to N-1
    • Iterate j = 1 to i+1
      • If dp[i - 1][j] is not -1, update dp[i][j] with dp[i - 1][j].
      • If dp[i - 1][j - 1] is non-negative before and after adding A[i], update dp[i][j] with the sum of dp[i - 1][j - 1] and A[i].
  • for i = N-1, iterate for [ j = 0 to N ] and return maximum ‘j’ for which dp[i][j] is non-negative.

03 Approach

Approach: We need to find the longest subsequence such that its prefix sum at each position is non-negative. Initially, the prefix sum is zero, and we start adding elements from left to right. If the current element can be taken without our prefix sum being negative, we take it irrespective of its value, as it increases the number of collected elements/balls.

At any position where the current element can’t be taken, try to maximize the prefix sum. Maintain a priority queue (min-heap) to store our collected elements in non-decreasing order. If we have collected any element ‘x’ that is more negative than the current element, then subtract ‘x’ from our prefix sum and assume that we have not collected that. In place of ‘x’, we will take the current element; although it doesn’t increase the count of elements, it makes the prefix sum more positive, which can be used in further processing.
 

Algorithm:

  • Create an empty priority queue and initialize the prefix sum with zero.
  • Iterate a for loop from i = 0 to N-1
    • Add A[i] to our prefix sum.
    • Assume A[i] is collected, so push it to the priority queue.
    • If prefix sum becomes negative
      • Pop the most negative element from the queue and subtract from the prefix sum as well.
  • Finally, our priority queue contains all collected elements. Return the number of elements in the queue.