CRStack.py

from collections import defaultdict
from typing import DefaultDict, Generic, Iterator, Optional, TypeVar

T = TypeVar("T")


class Node(Generic[T]):
    __slots__ = ("data", "bottom")

    def __init__(self, value : T, bottom: Optional["Node[T]"] = None) -> None:
        self.data = value
        self.bottom = bottom


class Stack(Generic[T]):
    __slots__ = ("head", "size", "frequency_map")

    def __init__(self) -> None:
        self.head: Optional[Node[T]] = None
        self.size: int = 0
        self.frequency_map: DefaultDict[T, int] = defaultdict(int)

    def push(self, value: T) -> None:
        self.head = Node(value, self.head)
        self.size += 1
        self.frequency_map[value] += 1

    def is_empty(self) -> bool:
        return self.head is None

    def pop(self) -> T:
        if self.head is None:
            raise IndexError("popping from an empty stack")

        node = self.head
        self.head = node.bottom
        self.size -= 1
        value = node.data
        self.frequency_map[value] -= 1

        if self.frequency_map[value] == 0:
            del self.frequency_map[value]

        return value

    def clear(self) -> None:
        self.head = None
        self.size = 0
        self.frequency_map.clear()

    def peek(self) -> T:
        if self.head is None:
            raise IndexError("peeking from an empty stack")

        return self.head.data

    def __iter__(self) -> Iterator[T]:
        node = self.head
        while node:
            yield node.data
            node = node.bottom

    def __contains__(self, item: T) -> bool:
        return self.frequency_map.get(item, 0) > 0

    def __len__(self) -> int:
        return self.size

    def __bool__(self) -> bool:
        return self.head is not None


def main() -> None:
    s = Stack()
    s.push(1)
    s.push(2)
    s.push(3)
    s.push(1)

    assert len(s) == 4

    assert 1 in s
    assert 2 in s
    assert 3 in s
    assert 4 not in s

    for i in s:
        print(i)

    s.pop()

    assert 1 in s
    assert 2 in s
    assert 3 in s
    assert s

    s.pop()

    assert 1 in s
    assert 2 in s
    assert 3 not in s

    s.pop()

    assert 1 in s
    assert 2 not in s
    assert 3 not in s

    s.pop()

    assert 0 not in s
    assert 1 not in s
    assert 2 not in s
    assert 3 not in s
    assert 4 not in s

    assert s.is_empty()
    assert len(s) == 0
    assert not s

    print("[OK]")


if __name__ == "__main__":
    main()