Skip to content

Commit e3c07f9

Browse files
authored
Add skew heap data structure. (#3238)
* Add skew heap data structure. * fixup! Add skew heap data structure. * fixup! Add skew heap data structure. * fixup! Add skew heap data structure. * Add tests. * Add __iter__ method. * fixup! Add __iter__ method.
1 parent 5fcd250 commit e3c07f9

File tree

1 file changed

+192
-0
lines changed

1 file changed

+192
-0
lines changed

Diff for: data_structures/heap/skew_heap.py

+192
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
#!/usr/bin/env python3
2+
3+
from __future__ import annotations
4+
5+
from typing import Generic, Iterable, Iterator, Optional, TypeVar
6+
7+
T = TypeVar("T")
8+
9+
10+
class SkewNode(Generic[T]):
11+
"""
12+
One node of the skew heap. Contains the value and references to
13+
two children.
14+
"""
15+
16+
def __init__(self, value: T) -> None:
17+
self._value: T = value
18+
self.left: Optional[SkewNode[T]] = None
19+
self.right: Optional[SkewNode[T]] = None
20+
21+
@property
22+
def value(self) -> T:
23+
"""Return the value of the node."""
24+
return self._value
25+
26+
@staticmethod
27+
def merge(
28+
root1: Optional[SkewNode[T]], root2: Optional[SkewNode[T]]
29+
) -> Optional[SkewNode[T]]:
30+
"""Merge 2 nodes together."""
31+
if not root1:
32+
return root2
33+
34+
if not root2:
35+
return root1
36+
37+
if root1.value > root2.value:
38+
root1, root2 = root2, root1
39+
40+
result = root1
41+
temp = root1.right
42+
result.right = root1.left
43+
result.left = SkewNode.merge(temp, root2)
44+
45+
return result
46+
47+
48+
class SkewHeap(Generic[T]):
49+
"""
50+
A data structure that allows inserting a new value and to pop the smallest
51+
values. Both operations take O(logN) time where N is the size of the
52+
structure.
53+
Wiki: https://en.wikipedia.org/wiki/Skew_heap
54+
Visualisation: https://www.cs.usfca.edu/~galles/visualization/SkewHeap.html
55+
56+
>>> list(SkewHeap([2, 3, 1, 5, 1, 7]))
57+
[1, 1, 2, 3, 5, 7]
58+
59+
>>> sh = SkewHeap()
60+
>>> sh.pop()
61+
Traceback (most recent call last):
62+
...
63+
IndexError: Can't get top element for the empty heap.
64+
65+
>>> sh.insert(1)
66+
>>> sh.insert(-1)
67+
>>> sh.insert(0)
68+
>>> list(sh)
69+
[-1, 0, 1]
70+
"""
71+
72+
def __init__(self, data: Optional[Iterable[T]] = ()) -> None:
73+
"""
74+
>>> sh = SkewHeap([3, 1, 3, 7])
75+
>>> list(sh)
76+
[1, 3, 3, 7]
77+
"""
78+
self._root: Optional[SkewNode[T]] = None
79+
for item in data:
80+
self.insert(item)
81+
82+
def __bool__(self) -> bool:
83+
"""
84+
Check if the heap is not empty.
85+
86+
>>> sh = SkewHeap()
87+
>>> bool(sh)
88+
False
89+
>>> sh.insert(1)
90+
>>> bool(sh)
91+
True
92+
>>> sh.clear()
93+
>>> bool(sh)
94+
False
95+
"""
96+
return self._root is not None
97+
98+
def __iter__(self) -> Iterator[T]:
99+
"""
100+
Returns sorted list containing all the values in the heap.
101+
102+
>>> sh = SkewHeap([3, 1, 3, 7])
103+
>>> list(sh)
104+
[1, 3, 3, 7]
105+
"""
106+
result = []
107+
while self:
108+
result.append(self.pop())
109+
110+
# Pushing items back to the heap not to clear it.
111+
for item in result:
112+
self.insert(item)
113+
114+
return iter(result)
115+
116+
def insert(self, value: T) -> None:
117+
"""
118+
Insert the value into the heap.
119+
120+
>>> sh = SkewHeap()
121+
>>> sh.insert(3)
122+
>>> sh.insert(1)
123+
>>> sh.insert(3)
124+
>>> sh.insert(7)
125+
>>> list(sh)
126+
[1, 3, 3, 7]
127+
"""
128+
self._root = SkewNode.merge(self._root, SkewNode(value))
129+
130+
def pop(self) -> T:
131+
"""
132+
Pop the smallest value from the heap and return it.
133+
134+
>>> sh = SkewHeap([3, 1, 3, 7])
135+
>>> sh.pop()
136+
1
137+
>>> sh.pop()
138+
3
139+
>>> sh.pop()
140+
3
141+
>>> sh.pop()
142+
7
143+
>>> sh.pop()
144+
Traceback (most recent call last):
145+
...
146+
IndexError: Can't get top element for the empty heap.
147+
"""
148+
result = self.top()
149+
self._root = SkewNode.merge(self._root.left, self._root.right)
150+
151+
return result
152+
153+
def top(self) -> T:
154+
"""
155+
Return the smallest value from the heap.
156+
157+
>>> sh = SkewHeap()
158+
>>> sh.insert(3)
159+
>>> sh.top()
160+
3
161+
>>> sh.insert(1)
162+
>>> sh.top()
163+
1
164+
>>> sh.insert(3)
165+
>>> sh.top()
166+
1
167+
>>> sh.insert(7)
168+
>>> sh.top()
169+
1
170+
"""
171+
if not self._root:
172+
raise IndexError("Can't get top element for the empty heap.")
173+
return self._root.value
174+
175+
def clear(self):
176+
"""
177+
Clear the heap.
178+
179+
>>> sh = SkewHeap([3, 1, 3, 7])
180+
>>> sh.clear()
181+
>>> sh.pop()
182+
Traceback (most recent call last):
183+
...
184+
IndexError: Can't get top element for the empty heap.
185+
"""
186+
self._root = None
187+
188+
189+
if __name__ == "__main__":
190+
import doctest
191+
192+
doctest.testmod()

0 commit comments

Comments
 (0)