Stacks (with "max" method)

1. Problem statement

Implement a stack with "max" method (Aziz et al., 2018, p. 106).

2. Insights

A stack supports two basic operations, push and pop, to add and remove elements from the stack in first-in-last-out (aka "FILO") order.

Keeping track of the max value is easy when we push in new elements into the stack. But when we pop an element out and it also happens to be the current max, how can we reconstruct the previous max value? We'd have to keep track of this somehow.

One way is to use a cache: for each element, we record what the current max value is. So we store both the element itself and the current max value on every push. Then when we pop this item out, we don't lose any information (we just look at the max value of the last-inserted item). This way we don't have to go through the entire stack (\(O(n)\) time complexity) to see which is the next max value.

3. Solutions

3.1. Store (elt, max)

code(1/2) 🔗
from linked_list.linked_list import LinkedList

class StackNaive:
    stack_naive_class_methods

We represent the stack with a linked list. We already implemented a linked list in "Linked lists", so we just reuse that here.

Aziz et al. (2018) use a native Python list to represent the stack. This is simpler to implement. But we choose the linked list instead because it more closely mirrors a true stack — namely, we are somewhat forced look at the last-inserted element only (unlike a list where one can do an index lookup to get any element in constant time).

stack_naive_class_methods(1/7) 🔗
def __init__(self, *args):
    self.ll = LinkedList()
    for elt in args:
        self.push(elt)

3.1.1. Insertion (push)

Insertion into a stack is done by adding an item to the end of the list. The linked list implementation we use inserts items by prepending to the linked list. This is already what we need.

What we need to add is the tracking of the max element whenever we push into the stack. So instead of storing just the element itself, we store a separate max value that represents the max value of all previously-pushed elements.

stack_naive_class_methods(2/7) 🔗
def push(self, elt: Any) -> None:
    if self.size() == 0:
        self.ll.insert((elt, elt))
    else:
        self.ll.insert((elt, max(elt, self.max())))
3.1.1.1. Complexity
  • Time: \(O(1)\).
  • Space: \(O(n)\), where \(n\) is the number of elements in the stack.

3.1.2. Max

For getting the max value, we just return the max variable that we saved along with the element.

stack_naive_class_methods(3/7) 🔗
def max(self):
    if self.size() == 0:
        return None
    node = self.ll.next
    (_, max) = node.elt
    return max

3.1.3. Deletion (pop)

Popping isn't really special at all — we just make sure to:

  1. delete the element, and
  2. return the value of the element to the caller.
stack_naive_class_methods(4/7) 🔗
def pop(self) -> Optional[Any]:
    node = self.ll.next
    (elt, _) = node.elt if node else (None, None)
    self.ll.delete_after()
    return elt
3.1.3.1. Complexity
  • Time: \(O(1)\).
  • Space: \(O(n)\), where \(n\) is the number of elements in the stack.

3.1.4. Print (__repr__)

This is useful for debugging. The code here follows the __repr__ example from CPython's datetime implementation.

stack_naive_class_methods(5/7) 🔗
def __repr__(self):
    elts = []
    node = self.ll
    while node.next:
        elts.append(repr(node.next.elt))
        node = node.next
        if node.next is None:
            break
    return "%s.%s(%s)" % (self.__class__.__module__,
                        self.__class__.__qualname__,
                        ", ".join(elts))

3.1.5. Size

stack_naive_class_methods(6/7) 🔗
def size(self) -> int:
    return self.ll.size()

3.1.6. Equality

stack_naive_class_methods(7/7) 🔗
def __eq__(self, other) -> bool:
    return self.ll == other.ll

3.2. Use two stacks

The downside of storing the max element with every element is that it can be wasteful. For example, if the stack has 100 elements of the same value, most of the max value information would be redundant.

We can deduplicate this redundancy by using an auxiliary stack, whose elements are of the form (max, count). The count simply stores how many times the current-highest max value is stored in the stack (so that popping it off would reduce this count value). The primary stack would simply store the elements as they are.

We inherit from the StackNaive class because we can reuse some of the existing code.

code(2/2) 🔗
class Stack(StackNaive):
    stack_class_methods

The main thing is that we have to add a new secondary stack (ll_max) to keep track of the maximum values.

stack_class_methods(1/11) 🔗
def __init__(self, *args):
    self.ll_max = LinkedList()
    super().__init__(*args)

3.2.1. Insertion (push)

Insertion is pretty simple: we first insert the element into the main stack, ll.

stack_class_methods(2/11) 🔗
def push(self, elt: Any) -> None:
    self.ll.insert(elt)

Now we have to update ll_max. The first edge case is if we are inserting the very first element into the main stack. If so, then by default this element is the current maximum, and so we insert it immediately. We set this element's counter to 1 as it is the first max item at this value.

stack_class_methods(3/11) 🔗
    if self.size() == 1:
        self.ll_max.insert((elt, 1))
        return

Now comes the case where we are inserting into a many-element stack. In this case we have to consider what is the current (leading) max value, and only insert into ll_max if the just-pushed element is at least the same size or greater the current maximum.

If the element is the same size as the current maximum, we just bump this maximum value's counter.

stack_class_methods(4/11) 🔗
    (max_cur, n) = self.ll_max.next.elt
    if elt == max_cur:
        self.ll_max.delete_after()
        self.ll_max.insert((max_cur, n+1))

Otherwise, we have to insert a new value into ll_max, just like we did for the case above for inserting into an empty stack.

stack_class_methods(5/11) 🔗
    elif elt > max_cur:
        self.ll_max.insert((elt, 1))

3.2.2. Max

Getting the max value is as easy as just observing the top of the ll_max stack.

stack_class_methods(6/11) 🔗
def max(self):
    if self.size() == 0:
        return None
    (max, _n) = self.ll_max.next.elt
    return max

3.2.3. Deletion (pop)

First we have to delete the element from the main stack.

stack_class_methods(7/11) 🔗
def pop(self) -> Optional[Any]:
    node = self.ll.next
    elt = node.elt if node else None
    self.ll.delete_after()

If we just popped the last element of the stack, then the stack is empty and there are no more maximum values, and so we also delete whatever last item remaining in ll_max.

stack_class_methods(8/11) 🔗
    if self.size() == 0:
        self.ll_max.delete_after()
        return elt

If the stack is not empty, then we have to check whether the just-popped element was equal to the current max value. That is, we can ignore the case where the element was smaller than the current max. E.g., if the stack is [3, 2, 1] with 3 as the first element, then the current max is 3 and popping either the 1 or the 2 will have no effect on the max. This is why we can ignore such cases.

It is also impossible for the element to be greater than the current max value, because of how we handle insertions (we update the max if the just-pushed element is greater than the current max).

So assuming that the just-popped element is equal to the max, we have to decrement the max value's counter in ll_max by 1. We do this by first deleting the max value, and restoring it to its counter with its decremented value, but only if it is above 1 (if zero, then we have to delete this max value outright).

stack_class_methods(9/11) 🔗
    (max_cur, n) = self.ll_max.next.elt
    if elt == max_cur:
        self.ll_max.delete_after()
        if n > 1:
            self.ll_max.insert((max_cur, n-1))

    return elt

3.2.4. Print (__repr__)

Printing is straightforward — we just reuse the __repr__() from the parent class StackNaive, but also marshal the second stack ll_max.

stack_class_methods(10/11) 🔗
def __repr__(self):
    elts = []
    node = self.ll_max
    while node.next:
        elts.append(repr(node.next.elt))
        node = node.next
        if node.next is None:
            break
    ll_repr = super().__repr__()
    ll_max_repr = "ll_max(%s)" % (", ".join(elts))
    return "%s (%s)" % (ll_repr, ll_max_repr)

3.2.5. Equality

Equality simply checks whether both the primary and secondary stacks are equal.

stack_class_methods(11/11) 🔗
def __eq__(self, other) -> bool:
    return self.ll == other.ll \
        and self.ll_max == other.ll_max

3.2.6. Complexity

  • Time: \(O(1)\).
  • Space: \(O(1)\) in the best case (when the maximum changes infrequently, such as the very first element being the max), \(O(n)\) in the worst case where every single newly-pushed element is greater than all previously-pushed elements (such that ll_max grows at the same pace as ll).

4. Tests

The from __future__ import annotations line is to enable self-referencing type annotations.

🔗
from __future__ import annotations
from hypothesis import given, strategies as st
import unittest
import random

from .stack import Stack, StackNaive

class Test(unittest.TestCase):
    test_cases

if __name__ == "__main__":
    unittest.main(exit=False)

4.1. Initialization

Here are some basic checks regarding initialization.

test_cases(1/6) 🔗
def test_init_empty_naive(self):
    stack = StackNaive()
    self.assertEqual(stack.size(), 0)

def test_init_empty(self):
    stack = Stack()
    self.assertEqual(stack.size(), 0)


def test_init_singleton_naive(self):
    stack = StackNaive(1)
    self.assertEqual(stack.size(), 1)

def test_init_singleton(self):
    stack = Stack(1)
    self.assertEqual(stack.size(), 1)


def test_init_multiple_naive(self):
    stack = StackNaive(1, 2, 3, 4, 5)
    self.assertEqual(stack.size(), 5)

def test_init_multiple(self):
    stack = Stack(1, 2, 3, 4, 5)
    self.assertEqual(stack.size(), 5)

4.2. Insertion (push)

test_cases(2/6) 🔗
def test_insertion_naive(self):
    stack = StackNaive()
    stack.push(1)
    stack.push(2)
    stack.push(3)
    self.assertEqual(stack, StackNaive(1, 2, 3))

def test_insertion(self):
    stack = Stack()
    stack.push(1)
    stack.push(2)
    stack.push(3)
    self.assertEqual(stack, Stack(1, 2, 3))


@given(st.lists(st.integers(min_value=1, max_value=100),
                min_size=0,
                max_size=20))
def test_insertion_random_naive(self, given_elts: list[int]):
    stack = StackNaive(*given_elts)
    self.assertEqual(stack.size(), len(given_elts))

@given(st.lists(st.integers(min_value=1, max_value=100),
                min_size=0,
                max_size=20))
def test_insertion_random(self, given_elts: list[int]):
    stack = Stack(*given_elts)
    self.assertEqual(stack.size(), len(given_elts))

4.3. Size

test_cases(3/6) 🔗
def test_size_naive(self):
    cases = [
        ([], 0),
        ([1], 1),
        ([1, 2], 2),
    ]
    for elts, expected in cases:
        stack = StackNaive(*elts)
        self.assertEqual(stack.size(), expected,
                         msg=f'{expected=}')

def test_size(self):
    cases = [
        ([], 0),
        ([1], 1),
        ([1, 2], 2),
    ]
    for elts, expected in cases:
        stack = Stack(*elts)
        self.assertEqual(stack.size(), expected,
                         msg=f'{expected=}')

4.4. Equality

test_cases(4/6) 🔗
def test_equality_naive(self):
    cases = [
        ([], [], True),
        ([1], [], False),
        ([1], [1], True),
        ([1], [1, 2], False),
    ]
    for xs, ys, expected in cases:
        x = StackNaive(*xs)
        y = StackNaive(*ys)
        self.assertEqual(x == y, expected,
                         msg=f'{x=} {y=}')

def test_equality(self):
    cases = [
        ([], [], True),
        ([1], [], False),
        ([1], [1], True),
        ([1], [1, 2], False),
    ]
    for xs, ys, expected in cases:
        x = Stack(*xs)
        y = Stack(*ys)
        self.assertEqual(x == y, expected,
                         msg=f'{x=} {y=}')

4.5. Deletion (pop)

test_cases(5/6) 🔗
def test_deletion_basic_naive(self):
    stack = StackNaive(1, 2, 3)
    stack.pop()
    self.assertEqual(stack, StackNaive(1, 2),
                    msg=f'{stack=}')

def test_deletion_basic(self):
    stack = Stack(1, 2, 3)
    stack.pop()
    self.assertEqual(stack, Stack(1, 2),
                    msg=f'{stack=}')


@given(st.lists(st.integers(min_value=1, max_value=100),
                min_size=1,
                max_size=20))
def test_deletion_random_naive(self, given_elts: list[int]):
    stack = StackNaive(*given_elts)
    deletions = random.randint(0, len(given_elts))
    for _ in range(deletions):
        stack.pop()
    self.assertEqual(stack.size(), len(given_elts) - deletions)

@given(st.lists(st.integers(min_value=1, max_value=100),
                min_size=1,
                max_size=20))
def test_deletion_random(self, given_elts: list[int]):
    stack = Stack(*given_elts)
    deletions = random.randint(0, len(given_elts))
    for _ in range(deletions):
        stack.pop()
    self.assertEqual(stack.size(), len(given_elts) - deletions)


def test_deletion_nop_naive(self):
    stack = StackNaive()
    stack.pop()
    self.assertEqual(stack.size(), 0)

def test_deletion_nop(self):
    stack = Stack()
    stack.pop()
    self.assertEqual(stack.size(), 0)

4.6. Max API

The stack should maintain the current max value as we push and pop elements with it.

test_cases(6/6) 🔗
def test_max_naive(self):
    cases = [
        ([1, 2, 3], [1, 2, 3], [2, 1]),
        ([2, 3, 10], [2, 3, 10], [3, 2]),
        ([10, 9, 8], [10, 10, 10], [10, 10]),
        ([1, 10, 5], [1, 10, 10], [10, 1]),
        ([2, 1, 2, 1], [2, 2, 2, 2], [2, 2, 2]),
        ([2, 1, 1, 1], [2, 2, 2, 2], [2, 2, 2]),
        ([2, 1, 8, 1], [2, 2, 8, 8], [8, 2, 2]),
        ([2, 1, 8, 8, 1], [2, 2, 8, 8, 8], [8, 8, 2, 2]),
    ]
    for elts, maxes, restored_maxes in cases:
        stack = StackNaive()
        for (elt, max) in zip(elts, maxes):
            stack.push(elt)
            self.assertEqual(stack.max(), max,
                             msg=f'{stack=}')
        for restored_max in restored_maxes:
            stack.pop()
            self.assertEqual(stack.max(), restored_max,
                             msg=f'{stack=}')

def test_max(self):
    cases = [
        ([1, 2, 3], [1, 2, 3], [2, 1]),
        ([2, 3, 10], [2, 3, 10], [3, 2]),
        ([10, 9, 8], [10, 10, 10], [10, 10]),
        ([1, 10, 5], [1, 10, 10], [10, 1]),
        ([2, 1, 2, 1], [2, 2, 2, 2], [2, 2, 2]),
        ([2, 1, 1, 1], [2, 2, 2, 2], [2, 2, 2]),
        ([2, 1, 8, 1], [2, 2, 8, 8], [8, 2, 2]),
        ([2, 1, 8, 8, 1], [2, 2, 8, 8, 8], [8, 8, 2, 2]),
    ]
    for elts, maxes, restored_maxes in cases:
        stack = Stack()
        for (elt, max) in zip(elts, maxes):
            stack.push(elt)
            self.assertEqual(stack.max(), max,
                             msg=f'{stack=}')
        for restored_max in restored_maxes:
            stack.pop()
            self.assertEqual(stack.max(), restored_max,
                             msg=f'{stack=}')

5. Export

Make this library available in other problems. 

from __future__ import annotations
from typing import Any, Optional
code

6. References

Aziz, A., Lee, T.-H., & Prakash, A. (2018). Elements of Programming Interviews in Python: The Insiders’ Guide. CreateSpace Independent Publishing Platform (25 July. 2018).