Add skew heap data structure.

data_structures/heap/skew_heap.py

@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+from typing import Generic, Iterable, Iterator, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+class SkewNode(Generic[T]):
+    """
+    One node of the skew heap. Contains the value and references to
+    two children.
+    """
+
+    def __init__(self, value: T) -> None:
+        self._value: T = value
+        self.left: Optional[SkewNode[T]] = None
+        self.right: Optional[SkewNode[T]] = None
+
+    @property
+    def value(self) -> T:
+        """Return the value of the node."""
+        return self._value
+
+    @staticmethod
+    def merge(
+        root1: Optional[SkewNode[T]], root2: Optional[SkewNode[T]]
+    ) -> Optional[SkewNode[T]]:
+        """Merge 2 nodes together."""
+        if not root1:
+            return root2
+
+        if not root2:
+            return root1
+
+        if root1.value > root2.value:
+            root1, root2 = root2, root1
+
+        result = root1
+        temp = root1.right
+        result.right = root1.left
+        result.left = SkewNode.merge(temp, root2)
+
+        return result
+
+
+class SkewHeap(Generic[T]):
+    """
+    A data structure that allows inserting a new value and to pop the smallest
+    values. Both operations take O(logN) time where N is the size of the
+    structure.
+    Wiki: https://en.wikipedia.org/wiki/Skew_heap
+    Visualisation: https://www.cs.usfca.edu/~galles/visualization/SkewHeap.html
+
+    >>> list(SkewHeap([2, 3, 1, 5, 1, 7]))
+    [1, 1, 2, 3, 5, 7]
+
+    >>> sh = SkewHeap()
+    >>> sh.pop()
+    Traceback (most recent call last):
+        ...
+    IndexError: Can't get top element for the empty heap.
+
+    >>> sh.insert(1)
+    >>> sh.insert(-1)
+    >>> sh.insert(0)
+    >>> list(sh)
+    [-1, 0, 1]
+    """
+
+    def __init__(self, data: Optional[Iterable[T]] = ()) -> None:
+        """
+        >>> sh = SkewHeap([3, 1, 3, 7])
+        >>> list(sh)
+        [1, 3, 3, 7]
+        """
+        self._root: Optional[SkewNode[T]] = None
+        for item in data:
+            self.insert(item)
+
+    def __bool__(self) -> bool:
+        """
+        Check if the heap is not empty.
+
+        >>> sh = SkewHeap()
+        >>> bool(sh)
+        False
+        >>> sh.insert(1)
+        >>> bool(sh)
+        True
+        >>> sh.clear()
+        >>> bool(sh)
+        False
+        """
+        return self._root is not None
+
+    def __iter__(self) -> Iterator[T]:
+        """
+        Returns sorted list containing all the values in the heap.
+
+        >>> sh = SkewHeap([3, 1, 3, 7])
+        >>> list(sh)
+        [1, 3, 3, 7]
+        """
+        result = []
+        while self:
+            result.append(self.pop())
+
+        # Pushing items back to the heap not to clear it.
+        for item in result:
+            self.insert(item)
+
+        return iter(result)
+
+    def insert(self, value: T) -> None:
+        """
+        Insert the value into the heap.
+
+        >>> sh = SkewHeap()
+        >>> sh.insert(3)
+        >>> sh.insert(1)
+        >>> sh.insert(3)
+        >>> sh.insert(7)
+        >>> list(sh)
+        [1, 3, 3, 7]
+        """
+        self._root = SkewNode.merge(self._root, SkewNode(value))
+
+    def pop(self) -> T:
+        """
+        Pop the smallest value from the heap and return it.
+
+        >>> sh = SkewHeap([3, 1, 3, 7])
+        >>> sh.pop()
+        1
+        >>> sh.pop()
+        3
+        >>> sh.pop()
+        3
+        >>> sh.pop()
+        7
+        >>> sh.pop()
+        Traceback (most recent call last):
+            ...
+        IndexError: Can't get top element for the empty heap.
+        """
+        result = self.top()
+        self._root = SkewNode.merge(self._root.left, self._root.right)
+
+        return result
+
+    def top(self) -> T:
+        """
+        Return the smallest value from the heap.
+
+        >>> sh = SkewHeap()
+        >>> sh.insert(3)
+        >>> sh.top()
+        3
+        >>> sh.insert(1)
+        >>> sh.top()
+        1
+        >>> sh.insert(3)
+        >>> sh.top()
+        1
+        >>> sh.insert(7)
+        >>> sh.top()
+        1
+        """
+        if not self._root:
+            raise IndexError("Can't get top element for the empty heap.")
+        return self._root.value
+
+    def clear(self):
+        """
+        Clear the heap.
+
+        >>> sh = SkewHeap([3, 1, 3, 7])
+        >>> sh.clear()
+        >>> sh.pop()
+        Traceback (most recent call last):
+            ...
+        IndexError: Can't get top element for the empty heap.
+        """
+        self._root = None
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()