randomized_heap.py

#!/usr/bin/env python3

from __future__ importannotations

importrandom
fromcollections.abcimportIterable
fromtypingimportAny, Generic, TypeVar

T=TypeVar("T", bound=bool)


classRandomizedHeapNode(Generic[T]):
"""
 One node of the randomized heap. Contains the value and references to
 two children.
 """

def__init__(self, value: T) ->None:
self._value: T=value
self.left: RandomizedHeapNode[T] |None=None
self.right: RandomizedHeapNode[T] |None=None

@property
defvalue(self) ->T:
"""
 Return the value of the node.

 >>> rhn = RandomizedHeapNode(10)
 >>> rhn.value
 10
 >>> rhn = RandomizedHeapNode(-10)
 >>> rhn.value
 -10
 """
returnself._value

@staticmethod
defmerge(
root1: RandomizedHeapNode[T] |None, root2: RandomizedHeapNode[T] |None
 ) ->RandomizedHeapNode[T] |None:
"""
 Merge 2 nodes together.

 >>> rhn1 = RandomizedHeapNode(10)
 >>> rhn2 = RandomizedHeapNode(20)
 >>> RandomizedHeapNode.merge(rhn1, rhn2).value
 10

 >>> rhn1 = RandomizedHeapNode(20)
 >>> rhn2 = RandomizedHeapNode(10)
 >>> RandomizedHeapNode.merge(rhn1, rhn2).value
 10

 >>> rhn1 = RandomizedHeapNode(5)
 >>> rhn2 = RandomizedHeapNode(0)
 >>> RandomizedHeapNode.merge(rhn1, rhn2).value
 0
 """
ifnotroot1:
returnroot2

ifnotroot2:
returnroot1

ifroot1.value>root2.value:
root1, root2=root2, root1

ifrandom.choice([True, False]):
root1.left, root1.right=root1.right, root1.left

root1.left=RandomizedHeapNode.merge(root1.left, root2)

returnroot1


classRandomizedHeap(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/Randomized_meldable_heap

 >>> RandomizedHeap([2, 3, 1, 5, 1, 7]).to_sorted_list()
 [1, 1, 2, 3, 5, 7]

 >>> rh = RandomizedHeap()
 >>> rh.pop()
 Traceback (most recent call last):
 ...
 IndexError: Can't get top element for the empty heap.

 >>> rh.insert(1)
 >>> rh.insert(-1)
 >>> rh.insert(0)
 >>> rh.to_sorted_list()
 [-1, 0, 1]
 """

def__init__(self, data: Iterable[T] |None= ()) ->None:
"""
 >>> rh = RandomizedHeap([3, 1, 3, 7])
 >>> rh.to_sorted_list()
 [1, 3, 3, 7]
 """
self._root: RandomizedHeapNode[T] |None=None

ifdata:
foritemindata:
self.insert(item)

definsert(self, value: T) ->None:
"""
 Insert the value into the heap.

 >>> rh = RandomizedHeap()
 >>> rh.insert(3)
 >>> rh.insert(1)
 >>> rh.insert(3)
 >>> rh.insert(7)
 >>> rh.to_sorted_list()
 [1, 3, 3, 7]
 """
self._root=RandomizedHeapNode.merge(self._root, RandomizedHeapNode(value))

defpop(self) ->T|None:
"""
 Pop the smallest value from the heap and return it.

 >>> rh = RandomizedHeap([3, 1, 3, 7])
 >>> rh.pop()
 1
 >>> rh.pop()
 3
 >>> rh.pop()
 3
 >>> rh.pop()
 7
 >>> rh.pop()
 Traceback (most recent call last):
 ...
 IndexError: Can't get top element for the empty heap.
 """

result=self.top()

ifself._rootisNone:
returnNone

self._root=RandomizedHeapNode.merge(self._root.left, self._root.right)

returnresult

deftop(self) ->T:
"""
 Return the smallest value from the heap.

 >>> rh = RandomizedHeap()
 >>> rh.insert(3)
 >>> rh.top()
 3
 >>> rh.insert(1)
 >>> rh.top()
 1
 >>> rh.insert(3)
 >>> rh.top()
 1
 >>> rh.insert(7)
 >>> rh.top()
 1
 """
ifnotself._root:
raiseIndexError("Can't get top element for the empty heap.")
returnself._root.value

defclear(self) ->None:
"""
 Clear the heap.

 >>> rh = RandomizedHeap([3, 1, 3, 7])
 >>> rh.clear()
 >>> rh.pop()
 Traceback (most recent call last):
 ...
 IndexError: Can't get top element for the empty heap.
 """
self._root=None

defto_sorted_list(self) ->list[Any]:
"""
 Returns sorted list containing all the values in the heap.

 >>> rh = RandomizedHeap([3, 1, 3, 7])
 >>> rh.to_sorted_list()
 [1, 3, 3, 7]
 """
result= []
whileself:
result.append(self.pop())

returnresult

def__bool__(self) ->bool:
"""
 Check if the heap is not empty.

 >>> rh = RandomizedHeap()
 >>> bool(rh)
 False
 >>> rh.insert(1)
 >>> bool(rh)
 True
 >>> rh.clear()
 >>> bool(rh)
 False
 """
returnself._rootisnotNone


if__name__=="__main__":
importdoctest

doctest.testmod()