Python/data_structures/heap/skew_heap.py
Suyash Dongre 05a5cdacc3
Added doctest to skew_heap.py ()
* Added doctest to skew_heap.py

* Update skew_heap.py

* Update data_structures/heap/skew_heap.py

Co-authored-by: Saptadeep Banerjee  <69459134+imSanko@users.noreply.github.com>

* Update skew_heap.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update skew_heap.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Saptadeep Banerjee <69459134+imSanko@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2024-01-18 13:39:27 +01:00

238 lines
5.5 KiB
Python

#!/usr/bin/env python3
from __future__ import annotations
from collections.abc import Iterable, Iterator
from typing import Any, Generic, TypeVar
T = TypeVar("T", bound=bool)
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: SkewNode[T] | None = None
self.right: SkewNode[T] | None = None
@property
def value(self) -> T:
"""
Return the value of the node.
>>> SkewNode(0).value
0
>>> SkewNode(3.14159).value
3.14159
>>> SkewNode("hello").value
'hello'
>>> SkewNode(None).value
>>> SkewNode(True).value
True
>>> SkewNode([]).value
[]
>>> SkewNode({}).value
{}
>>> SkewNode(set()).value
set()
>>> SkewNode(0.0).value
0.0
>>> SkewNode(-1e-10).value
-1e-10
>>> SkewNode(10).value
10
>>> SkewNode(-10.5).value
-10.5
>>> SkewNode().value
Traceback (most recent call last):
...
TypeError: SkewNode.__init__() missing 1 required positional argument: 'value'
"""
return self._value
@staticmethod
def merge(
root1: SkewNode[T] | None, root2: SkewNode[T] | None
) -> SkewNode[T] | None:
"""
Merge 2 nodes together.
>>> SkewNode.merge(SkewNode(10),SkewNode(-10.5)).value
-10.5
>>> SkewNode.merge(SkewNode(10),SkewNode(10.5)).value
10
>>> SkewNode.merge(SkewNode(10),SkewNode(10)).value
10
>>> SkewNode.merge(SkewNode(-100),SkewNode(-10.5)).value
-100
"""
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
Visualization: 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: Iterable[T] | None = ()) -> None:
"""
>>> sh = SkewHeap([3, 1, 3, 7])
>>> list(sh)
[1, 3, 3, 7]
"""
self._root: SkewNode[T] | None = None
if data:
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: list[Any] = []
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 | None:
"""
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) if self._root else None
)
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) -> None:
"""
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()