Greedy: Find maximum sum possible equal sum of three stacks

Problem Statement:

You are given 3 arrays that represents a stack.

You need to remove the top elements of the stack such that the sum of all the stacks are equal.

Example

stack1[ ] = {4, 3, 2}
stack2[ ] = {3, 2, 1, 1, 1}
stack3[ ] = {1, 1, 4, 1}

Output: 5

Solution

We need to have a greedy approach.

Take 3 variables that will hold the sum of all the 3 arrays

If all the 3 values are same, then it is the max sum

Else remove the top element of the stack with maximum sum among the three.

Solution in C++

#include<iostream>
#include<vector>
#include<algorithm>
#include<queue>

using namespace std;

int get_max_sum(vector<int>& stack1, vector<int> &stack2, vector<int> &stack3, 
                int n1, int n2, int n3)
{
    int sum1 = 0, sum2 = 0, sum3 = 0;
 
    for (int i = 0; i < n1; i++)
        sum1 += stack1[i];
 
    for (int i = 0; i < n2; i++)
        sum2 += stack2[i];
 
    for (int i = 0; i < n3; i++)
        sum3 += stack3[i];
 
    int top1 = 0, top2 = 0, top3 = 0;
    while (1) {
        if (top1 == n1 || top2 == n2 || top3 == n3)
            return 0;
 
        if (sum1 == sum2 && sum2 == sum3)
            return sum1;
        if (sum1 >= sum2 && sum1 >= sum3)
            sum1 -= stack1[top1++];
        else if (sum2 >= sum1 && sum2 >= sum3)
            sum2 -= stack2[top2++];
        else if (sum3 >= sum2 && sum3 >= sum1)
            sum3 -= stack3[top3++];
    }
}
 
// Driven Program
int main()
{
    vector<int> stack1 = { 3, 2, 1, 1, 1 };
    vector<int> stack2 = { 4, 3, 2 };
    vector<int> stack3 = { 1, 1, 4, 1 };
 
    int n1 = stack1.size();
    int n2 = stack2.size();
    int n3 = stack3.size();
 
    cout << "The max sum is " <<get_max_sum(stack1, stack2, stack3, n1, n2, n3)
         << endl;
    return 0;
}

Output:

The max sum is 5