Search code examples
pythondata-structuresstack

Implementing a Stack in Python: Problem with the POP function


I'm trying to implement the Stack Data Structure in Python from scratch. Unfortunately, I have a problem with the pop functionality.

The pop function should delete the top element of the stack. Unfortunately, in my implementation, running pop and then printing the array returns the same array as before the pop function.

How to better implement the pop function? I could do self.top = self.top.prev, but I am not sure if in the Stack Data Structure we have access to previous elements.

# FOLLOWS LIFO


class StackElement:
    def __init__(self, value, prev=None, next=None):
        self.value = value
        self.prev = prev # Not sure if this is available in the Stack DS...
        self.next = next


class Stack:
    def __init__(self, values):
        self.bottom = None
        self.top = None

        for v in values:
            self.push(v)

    def __str__(self):
        res = [str(x.value) for x in self]
        return " -> ".join(res)

    def __iter__(self):
        current = self.bottom
        while current:
            yield current
            current = current.next

    def __len__(self):
        res = 0
        current = self.bottom
        while current:
            res += 1
            current = current.next
        return res

    def push(self, value):
        if self.bottom is None:
            self.bottom = self.top = StackElement(value)
        else:
            self.top.next = StackElement(value)
            self.top = self.top.next
        return self.top

    def pop(self):
        if not self.is_empty():
            self.top = None
        return self.top

    def peek(self):
        if not self.is_empty():
            return self.top

    def is_empty(self):
        if len(self) != 0:
            return False
        return True


if __name__ == "__main__":
    s = Stack(["2", "3", "4"])
    print(s)
    s.push("5")
    print(s, s.top.value, s.bottom.value)
    s.pop()
    print(s, s.top, s.bottom.value)
    res = [str(x.value) for x in s]
    print(res)
    print(s)

Console Output:

2 -> 3 -> 4
2 -> 3 -> 4 -> 5 5 2
2 -> 3 -> 4 -> 5 None 2
['2', '3', '4', '5']
2 -> 3 -> 4 -> 5

Solution

  • The assignment self.top = None in your pop method will not achieve what you want. It will just update that top attribute, but the node that precedes it will still link to that element via its next attribute. You really need to break that link by setting the appropriate next attribute to None. Also, you don't want to set top to None just like that. After the pop action there might still be more elements below that top, and the top attribute should be updated to reference that element that now has become the top of the stack.

    Besides this issue there are some other remarks to make:

    • In StackElement, don't name that parameter next, as that is the name of a native Python function. Use some other name, like nxt. For naming an attribute there is no such confusion possible, so that self.next can stay.

    • You have implemented a doubly linked list, but actually you only need a singly linked list. As you only pop from the "top" side, you only need prev, not next.

    • We could also omit bottom for a pure stack, but then we need to reconsider your iteration direction. Iterating over a stack is not really a plain stack feature, but if you do want to have it, it makes sense to iterate from top to bottom, and not from bottom to top. Then you can leave out bottom as well.

    • In __len__ you iterate the whole stack, which will work, but if you add an attribute to your stack that keeps track of the current length, then you can avoid this loop.

    • Your methods should not return StackElement instances. This should be considered private information. Only provide an interface with values that are stored in the stack, never the nodes. This also applies to __iter__. To indicate that the top attribute is private, prefix it with underscore.

    • The push method should not have to return anything (just None).

    • The StackElement has prev as optional argument. Using this argument in your push method can greatly reduce code.

    • is_empty can simply verify the length or the value of the top attribute and turn that into a boolean. No need for an if statement here.

    Here is your code with the above remarks taken into account:

    class StackElement:
        # Don't name argument "next", but you actually don't need it
        def __init__(self, value, prev=None):
            self.value = value
            self.prev = prev # No need for next
    
    class Stack:
        def __init__(self, values):
            # Use underscore to indicate these are considered private
            self._top = None  # No need for bottom
            self._len = 0  # Maintain length to avoid counting in __len__
    
            for v in values:
                self.push(v)
    
        def __str__(self):
            # Don't create a list. The iterator is enough for join
            # We reverse the arrow, because the iteration starts at the top
            return " <- ".join(map(str, self))
    
        def __iter__(self):
            current = self._top  # Start at top to avoid the bottom ref.
            while current:
                yield current.value  # Don't expose nodes, only values
                current = current.prev
    
        def __len__(self):
            return self._len
    
        def push(self, value):
            # Provide the prev argument to simplify code:
            self._top = StackElement(value, self._top)
            self._len += 1 # Keep length updated
            # Push should not return anything
    
        def pop(self):
            if not self.is_empty():
                value = self._top.value
                self._top = self._top.prev
                self._len -= 1  # Keep length updated
                return value  # Return value, not a node
    
        def peek(self):
            if not self.is_empty():
                return self._top.value # Return value, not a node
    
        def is_empty(self):
            return not self._top # no if statement needed