Design an algorithm that can return the Maximum item of a stack in O(1) running time complexity. We can use O(N) extra memory! : Stack Again : Chapter 3 : In Python

Today we will try to build a solution of a Stack problem.

We will build the stack from scratch and try to write a helper method for the stack which can fetch the Maximum element of the Stack.

So we are allowed to use O(N) extra memory but we have to fetch the item in constant Time complexity of O(1).

If you are reading this article for first time you can refer to my earlier article on Stack in python as below -

Creating a Stack in python : Data Structure : Stack : Chapter 2 : In Python | by Sanjit Khasnobis | Apr, 2022 | Medium

Now without further delay let us first analyze the scenario. So we have to find the Maximum Number from Stack with O(1) time complexity.

So, that does mean we have to get a stack in a form where the top element is the Maximum Element. So in that case if we just return the top element using our pop() method we can easily achieve O(1) Time Complexity.

So we will try to apply below process to build the final stack -

Step 1: We will define 2 stacks , 1 as main_stack and another as max_stack

Step 2: At the 1st insertion, We will insert or Push new element to main_stack as well as in max_stack.

Step 3: Now, from next element we will push the element as is in main_stack. But we will compare the new element with the existing last inserted element of max_stack(here the first one). If the current element is greater than the previous element inserted at max_stack then we will push this greater element to max_stack, if not we will discard this element and insert the previous element again in the max_stack.

So at any point of time, only the greater element is inserted in max_stack.

I am trying to explain the algorithm using the picture below.

Algorithm to Find Max of Stack

So the first element arrived 35.

We put that in both the stacks as-is as it is the first element of the stack.

Now the next element is 23.

We put 23 as-is at main_stack. But we compare 23 with the last element inserted in max_stack. As 23 < 35 , so we copy the last max element from max_stack and insert it to next index at max_stack.

Now 46 arrived. As 46 > 35 , so we now insert 46 at both main_stack and max_stack.

Assume the stack is finished we want to retrieve the maximum element of stack.

If you observe very closely by now we have created max_stack whose top element is the maximum number in whole stack.

So just doing a single pop() operation will return me the highest or Maximum number in stack. As it is a pop() operation so the Time Complexity will be always O(1), be it a stack of 3 element or be it a stack of 3 billion elements.

And we have used the same stack size as of main_stack so we run into O(N) space complexity as agreed at problem statement.

So, Algorithm part is finished, I will encourage everyone to try by own to write the corresponding code.

But I am sharing the code I try to build. Lets have fun.

class MaxStack():

    def __init__(self):
        self.main_stack = []
        self.max_stack  = []

    def push(self,data):
        self.main_stack.append(data)

        if(len(self.main_stack) == 1):
            self.max_stack.append(data)

        if(data > self.max_stack[-1]):
            self.max_stack.append(data)
        else:
            self.max_stack.append(self.max_stack[-1])


    def pop(self):
        self.max_stack.pop()
        return self.main_stack.pop()

    def getMax(self):
        return self.max_stack[-1]

if __name__ == "__main__":
    maxstack = MaxStack()
    maxstack.push(1234)
    maxstack.push(1234)
    maxstack.push(23457)
    maxstack.push(1234)
    maxstack.push(2345)
    maxstack.push(4567)
    maxstack.push(6789)
    maxstack.push(9001)
    maxstack.push(9002899)
    maxstack.push(9003)
    maxstack.push(9004)
    maxstack.push(9005)

    print("Maximum Number in Stack ::" + str(maxstack.getMax()))

At first we have defined 2 separate stacks at constructor of the class as below -

def __init__(self):
        self.main_stack = []
        self.max_stack  = []

Now push() method is as discussed.

The first element of main_stack will be as-is copied to max_stack.

Now the next block will check if the next coming element is greater than the existing max_stack element , if yes it will append the data to max_stack if not the previous Max element will be copied and again pushed to max_stack.

def push(self,data):
        self.main_stack.append(data)

        if(len(self.main_stack) == 1):
            self.max_stack.append(data)

        if(data > self.max_stack[-1]):
            self.max_stack.append(data)
        else:
            self.max_stack.append(self.max_stack[-1])

pop() method need to pop() both main_stack as well max_stack as below.

def pop(self):
        self.max_stack.pop()
        return self.main_stack.pop()

Now the final method is getMax(). getMax() method will return the top element from the stack which is the Maximum number in the stack.

def getMax(self):
        return self.max_stack[-1]

When we tested the code we get following output as below -

9002899 is the Maximum Number in this Stack

The getMax() function fetches the Maximum number as 9002899.

Hope you have enjoyed the article for Stack.

Please subscribe to read more about data engineering, data structure, azure data cloud, Aws Data Cloud, python , scala, Data Vault , Data Warehouse , AI and ML.

Happy Reading!! Happy Coding!!