How to implement 3 stacks with one array?

Question

Sometimes, I come across the following interview question: How to implement 3 stacks with one array ? Of course, any static allocation is not a solution.

Answer 1

Dr. belisarius's answer explains the basic algorithm, but doesn't go in the details, and as we know, the devil is always in the details. I coded up a solution in Python 3, with some explanation and a unit test. All operations run in constant time, as they should for a stack.

# One obvious solution is given array size n, divide up n into 3 parts, and allocate floor(n / 3) cells
# to two stacks at either end of the array, and remaining cells to the one in the middle. This strategy is not
# space efficient because even though there may be space in the array, one of the stack may overflow.
#
# A better approach is to have two stacks at either end of the array, the left one growing on the right, and the
# right one growing on the left. The middle one starts at index floor(n / 2), and grows at both ends. When the
# middle stack size is even, it grows on the right, and when it's odd, it grows on the left. This way, the middle
# stack grows evenly and minimizes the changes of overflowing one of the stack at either end.
#
# The rest is pointer arithmetic, adjusting tops of the stacks on push and pop operations.

class ThreeStacks:
    def __init__(self, n: int):
        self._arr: List[int] = [0] * n
        self._tops: List[int] = [-1, n, n // 2]
        self._sizes: List[int] = [0] * 3
        self._n = n

    def _is_stack_3_even_size(self):
        return self._sizes[2] % 2 == 0

    def _is_stack_3_odd_size(self):
        return not self._is_stack_3_even_size()

    def is_empty(self, stack_number: int) -> bool:
        return self._sizes[stack_number] == 0

    def is_full(self, stack_number: int) -> bool:
        if stack_number == 0 and self._is_stack_3_odd_size():
            return self._tops[stack_number] == self._tops[2] - self._sizes[2]
        elif stack_number == 1 and self._is_stack_3_even_size():
            return self._tops[stack_number] == self._tops[2] + self._sizes[2]

        return (self._is_stack_3_odd_size() and self._tops[0] == self._tops[2] - self._sizes[2]) or \
               (self._is_stack_3_even_size() and self._tops[1] == self._tops[2] + self._sizes[2])

    def pop(self, stack_number: int) -> int:
        if self.is_empty(stack_number):
            raise RuntimeError(f"Stack : {stack_number} is empty")

        x: int = self._arr[self._tops[stack_number]]
        if stack_number == 0:
            self._tops[stack_number] -= 1
        elif stack_number == 1:
            self._tops[stack_number] += 1
        else:
            if self._is_stack_3_even_size():
                self._tops[stack_number] += (self._sizes[stack_number] - 1)
            else:
                self._tops[stack_number] -= (self._sizes[stack_number] - 1)

        self._sizes[stack_number] -= 1
        return x

    def push(self, item: int, stack_number: int) -> None:
        if self.is_full(stack_number):
            raise RuntimeError(f"Stack: {stack_number} is full")

        if stack_number == 0:
            self._tops[stack_number] += 1
        elif stack_number == 1:
            self._tops[stack_number] -= 1
        else:
            if self._is_stack_3_even_size():
                self._tops[stack_number] += self._sizes[stack_number]
            else:
                self._tops[stack_number] -= self._sizes[stack_number]
        self._arr[self._tops[stack_number]] = item
        self._sizes[stack_number] += 1

    def __repr__(self):
        return str(self._arr)

Test:

def test_stack(self):
    stack = ThreeStacks(10)

    for i in range(3):
        with pytest.raises(RuntimeError):
            stack.pop(i)

    for i in range(1, 4):
        stack.push(i, 0)

    for i in range(4, 7):
        stack.push(i, 1)

    for i in range(7, 11):
        stack.push(i, 2)

    for i in range(3):
        with pytest.raises(RuntimeError):
            stack.push(1, i)

    assert [stack.pop(i) for i in range(3)] == [3, 6, 10]
    assert [stack.pop(i) for i in range(3)] == [2, 5, 9]
    assert [stack.pop(i) for i in range(3)] == [1, 4, 8]

    for i in range(2):
        assert stack.is_empty(i)

    assert not stack.is_empty(2)
    assert stack.pop(2) == 7
    assert stack.is_empty(2)