diff --git a/data_structures/heap/skew_heap.py b/data_structures/heap/skew_heap.py new file mode 100644 index 000000000..417a383f7 --- /dev/null +++ b/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()