change KeyObject.key type to float (to be able to safely assign math.inf which is float), add asserts to check if not None, many improvements to type hinting to make mypy happy

Author: wojtask

src/book/chapter6/problem1.py

@@ -1,10 +1,11 @@
 from book.chapter6.section5 import max_heap_insert
 from book.data_structures import Heap
 from book.data_structures import KeyObject
+from book.data_structures import T
 from util import range_of
 
 
-def build_max_heap_(A: Heap[KeyObject], n: int) -> None:
+def build_max_heap_(A: Heap[KeyObject[T]], n: int) -> None:
     """Converts an array into a max-heap by repeatedly inserting elements to the heap.
 
     Implements:

src/book/chapter6/section5.py

@@ -4,10 +4,11 @@
 from book.chapter6.section2 import max_heapify
 from book.data_structures import KeyObject
 from book.data_structures import PriorityQueue
+from book.data_structures import T
 from solutions.chapter2.section1.exercise4 import linear_search
 
 
-def max_heap_maximum(A: PriorityQueue) -> KeyObject:
+def max_heap_maximum(A: PriorityQueue[T]) -> KeyObject[T]:
     """Returns the element of the dynamic set with the largest key.
 
     Implements:
@@ -24,7 +25,7 @@ def max_heap_maximum(A: PriorityQueue) -> KeyObject:
     return A[1]
 
 
-def max_heap_extract_max(A: PriorityQueue) -> KeyObject:
+def max_heap_extract_max(A: PriorityQueue[T]) -> KeyObject[T]:
     """Removes and returns the element of the dynamic set with the largest key.
 
     Implements:
@@ -43,7 +44,7 @@ def max_heap_extract_max(A: PriorityQueue) -> KeyObject:
     return max
 
 
-def max_heap_increase_key(A: PriorityQueue, x: KeyObject, k: int) -> None:
+def max_heap_increase_key(A: PriorityQueue[T], x: KeyObject[T], k: float) -> None:
     """Increases the value of the element's key to the new value.
 
     Implements:
@@ -63,11 +64,13 @@ def max_heap_increase_key(A: PriorityQueue, x: KeyObject, k: int) -> None:
         i = parent(i)
 
 
-def __find_index_of_object(A: PriorityQueue, x: KeyObject) -> int:
-    return linear_search(A, A.heap_size, x)
+def __find_index_of_object(A: PriorityQueue[T], x: KeyObject[T]) -> int:
+    index = linear_search(A, A.heap_size, x)
+    assert index is not None
+    return index
 
 
-def max_heap_insert(A: PriorityQueue, x: KeyObject, n: int) -> None:
+def max_heap_insert(A: PriorityQueue[T], x: KeyObject[T], n: int) -> None:
     """Inserts a new element into the dynamic set.
 
     Implements:

src/book/data_structures.py

@@ -3,14 +3,12 @@
 from builtins import len
 from collections.abc import Callable
 from typing import Any
-from typing import Dict
 from typing import Generic
 from typing import Literal
 from typing import TypeVar
 
 from typing_extensions import Protocol
-
-from util import range_of
+from typing_extensions import TypeAlias
 
 T = TypeVar('T')
 CT = TypeVar('CT', bound='Comparable')
@@ -35,7 +33,7 @@ class Indexable(Protocol[T]):
     __setitem__: Callable[[int, T], None]
 
 
-Bit = Literal[0, 1]
+Bit: TypeAlias = Literal[0, 1]
 
 
 class Array(Generic[T]):
@@ -140,35 +138,35 @@ class Heap(Array[CT]):
     heap_size: int = 0
 
 
-class KeyObject:
-    def __init__(self, key: int, data=None) -> None:
+class KeyObject(Generic[T]):
+    def __init__(self, key: float, data: T) -> None:
         self.key = key
         self.data = data
 
-    def __lt__(self, other: KeyObject) -> bool:
+    def __lt__(self, other: KeyObject[T]) -> bool:
         return self.key < other.key
 
-    def __le__(self, other: KeyObject) -> bool:
+    def __le__(self, other: KeyObject[T]) -> bool:
         return self.key <= other.key
 
-    def __gt__(self, other: KeyObject) -> bool:
+    def __gt__(self, other: KeyObject[T]) -> bool:
         return self.key > other.key
 
-    def __ge__(self, other: KeyObject) -> bool:
+    def __ge__(self, other: KeyObject[T]) -> bool:
         return self.key >= other.key
 
 
 # TODO(#476) Consider removing this data structure, and instead implement the priority queue directly on a heap.
 #   Clarify if we need the mapping between objects and heap indices, because as mentioned in the book, a heap
 #   implementing a priority queue contains pointers to the objects, which should suffice to address the objects.
-class PriorityQueue(Heap[KeyObject]):
+class PriorityQueue(Heap[KeyObject[T]]):
 
     def __init__(self, n: int) -> None:
         # TODO(#21) implement the object mappings with a hash table instead of Python dictionary
         super().__init__(1, n)
         self.heap_size: int = 0
-        self.mapping: Dict[KeyObject, int] = {}
+        self.mapping: dict[KeyObject[T], int] = {}
 
-    def __setitem__(self, index: int, value: KeyObject) -> None:
+    def __setitem__(self, index: int, value: KeyObject[T]) -> None:
         super().__setitem__(index, value)
         self.mapping[value] = index

src/solutions/chapter6/section5/exercise10.py

@@ -4,9 +4,10 @@
 from book.chapter6.section5 import max_heap_increase_key
 from book.data_structures import KeyObject
 from book.data_structures import PriorityQueue
+from book.data_structures import T
 
 
-def max_heap_delete(A: PriorityQueue, x: KeyObject) -> None:
+def max_heap_delete(A: PriorityQueue[T], x: KeyObject[T]) -> None:
     """Deletes an element from a max-heap.
 
     Implements:

src/solutions/chapter6/section5/exercise11.py

@@ -1,27 +1,29 @@
 from __future__ import annotations
 
-from typing import Generic
+from typing_extensions import TypeAlias
 
 from book.data_structures import Array
 from book.data_structures import KeyObject
 from book.data_structures import PriorityQueue
-from book.data_structures import T
 from solutions.chapter6.section5.exercise3 import min_heap_extract_min
 from solutions.chapter6.section5.exercise3 import min_heap_insert
 from util import range_of
 
 
 # TODO(#20) replace by the general list node and list classes once they will have been added
-class SortedListNode(Generic[T]):
-    key: T
-    next: SortedListNode[T] | None = None
+class SortedListNode:
+    key: float
+    next: SortedListNode | None = None
 
 
-class SortedList(Generic[T]):
-    head: SortedListNode[T] | None = None
+class SortedList:
+    head: SortedListNode | None = None
 
 
-def merge_sorted_lists(sorted_lists: Array[SortedList[T]], k: int) -> SortedList[T]:
+NodePayloadType: TypeAlias = tuple[SortedListNode, SortedList]
+
+
+def merge_sorted_lists(sorted_lists: Array[SortedList], k: int) -> SortedList:
     """Merges sorted lists into one sorted list.
 
     Args:
@@ -31,24 +33,26 @@ def merge_sorted_lists(sorted_lists: Array[SortedList[T]], k: int) -> SortedList
     Returns:
         The sorted list containing all elements from the input sorted lists.
     """
-    Q = PriorityQueue(k)
+    Q = PriorityQueue[NodePayloadType](k)
     for i in range_of(1, to=k):
         sorted_list = sorted_lists[i]
         if sorted_list.head is not None:
             x = sorted_list.head
             sorted_list.head = sorted_list.head.next
             x.next = None
-            min_heap_insert(Q, KeyObject(key=x.key, data=(x, sorted_list)), k)
-    merged_list = SortedList[T]()
-    tail = None
+            min_heap_insert(Q, KeyObject[NodePayloadType](key=x.key, data=(x, sorted_list)), k)
+    merged_list = SortedList()
+    tail: SortedListNode | None = None
     while Q.heap_size > 0:
-        element, list_ = min_heap_extract_min(Q).data
+        head_object = min_heap_extract_min(Q)
+        element, list_ = head_object.data
         if merged_list.head is None:
             tail = merged_list.head = element
         else:
+            assert tail is not None
             tail.next = element
             tail = tail.next
         if list_.head is not None:
-            min_heap_insert(Q, KeyObject(key=list_.head.key, data=(list_.head, list_)), k)
+            min_heap_insert(Q, KeyObject[NodePayloadType](key=list_.head.key, data=(list_.head, list_)), k)
             list_.head = list_.head.next
     return merged_list

src/solutions/chapter6/section5/exercise3.py

@@ -3,11 +3,12 @@
 from book.chapter6.section1 import parent
 from book.data_structures import KeyObject
 from book.data_structures import PriorityQueue
+from book.data_structures import T
 from solutions.chapter2.section1.exercise4 import linear_search
 from solutions.chapter6.section2.exercise3 import min_heapify
 
 
-def min_heap_minimum(A: PriorityQueue) -> KeyObject:
+def min_heap_minimum(A: PriorityQueue[T]) -> KeyObject[T]:
     """Returns the element of the dynamic set with the smallest key.
 
     Implements:
@@ -24,7 +25,7 @@ def min_heap_minimum(A: PriorityQueue) -> KeyObject:
     return A[1]
 
 
-def min_heap_extract_min(A: PriorityQueue) -> KeyObject:
+def min_heap_extract_min(A: PriorityQueue[T]) -> KeyObject[T]:
     """Removes and returns the element of the dynamic set with the smallest key.
 
     Implements:
@@ -43,7 +44,7 @@ def min_heap_extract_min(A: PriorityQueue) -> KeyObject:
     return min
 
 
-def min_heap_decrease_key(A: PriorityQueue, x: KeyObject, k: int) -> None:
+def min_heap_decrease_key(A: PriorityQueue[T], x: KeyObject[T], k: float) -> None:
     """Decreases the value of the element's key to the new value.
 
     Implements:
@@ -63,11 +64,13 @@ def min_heap_decrease_key(A: PriorityQueue, x: KeyObject, k: int) -> None:
         i = parent(i)
 
 
-def __find_index_of_object(A: PriorityQueue, x: KeyObject) -> int:
-    return linear_search(A, A.heap_size, x)
+def __find_index_of_object(A: PriorityQueue[T], x: KeyObject[T]) -> int:
+    index = linear_search(A, A.heap_size, x)
+    assert index is not None
+    return index
 
 
-def min_heap_insert(A: PriorityQueue, x: KeyObject, n: int) -> None:
+def min_heap_insert(A: PriorityQueue[T], x: KeyObject[T], n: int) -> None:
     """Inserts a new element into the dynamic set.
 
     Implements:

src/solutions/chapter6/section5/exercise4.py

@@ -1,10 +1,11 @@
 from book.chapter6.section2 import max_heapify
 from book.data_structures import KeyObject
 from book.data_structures import PriorityQueue
+from book.data_structures import T
 from solutions.chapter2.section1.exercise4 import linear_search
 
 
-def max_heap_decrease_key(A: PriorityQueue, x: KeyObject, k: int) -> None:
+def max_heap_decrease_key(A: PriorityQueue[T], x: KeyObject[T], k: float) -> None:
     """Increases the value of the element's key to the new value.
 
     Implements:
@@ -22,5 +23,7 @@ def max_heap_decrease_key(A: PriorityQueue, x: KeyObject, k: int) -> None:
     max_heapify(A, i)
 
 
-def __find_index_of_object(A: PriorityQueue, x: KeyObject) -> int:
-    return linear_search(A, A.heap_size, x)
+def __find_index_of_object(A: PriorityQueue[T], x: KeyObject[T]) -> int:
+    index = linear_search(A, A.heap_size, x)
+    assert index is not None
+    return index

src/solutions/chapter6/section5/exercise8.py

@@ -1,10 +1,11 @@
 from book.chapter6.section1 import parent
 from book.data_structures import KeyObject
 from book.data_structures import PriorityQueue
+from book.data_structures import T
 from solutions.chapter2.section1.exercise4 import linear_search
 
 
-def max_heap_increase_key_(A: PriorityQueue, x: KeyObject, k: int) -> None:
+def max_heap_increase_key_(A: PriorityQueue[T], x: KeyObject[T], k: float) -> None:
     """Increases the value of the element's key to the new value. This version of the procedure uses fewer object
     assignments than the original Max-Heap-Insert-Key.
 
@@ -26,5 +27,7 @@ def max_heap_increase_key_(A: PriorityQueue, x: KeyObject, k: int) -> None:
     A[i] = x
 
 
-def __find_index_of_object(A: PriorityQueue, x: KeyObject) -> int:
-    return linear_search(A, A.heap_size, x)
+def __find_index_of_object(A: PriorityQueue[T], x: KeyObject[T]) -> int:
+    index = linear_search(A, A.heap_size, x)
+    assert index is not None
+    return index

src/solutions/chapter6/section5/exercise9.py

@@ -2,15 +2,16 @@
 from book.chapter6.section5 import max_heap_insert
 from book.data_structures import KeyObject
 from book.data_structures import PriorityQueue
+from book.data_structures import T
 from solutions.chapter6.section5.exercise3 import min_heap_extract_min
 from solutions.chapter6.section5.exercise3 import min_heap_insert
 
 
-class ControlledPriorityQueue(PriorityQueue):
+class ControlledPriorityQueue(PriorityQueue[T]):
     priority: int = 1
 
 
-def min_heap_enqueue(A: ControlledPriorityQueue, x: KeyObject, n: int) -> None:
+def min_heap_enqueue(A: ControlledPriorityQueue[T], x: KeyObject[T], n: int) -> None:
     """Inserts an element into a first-in, first-out queue implemented with a priority queue.
 
     Implements:
@@ -26,7 +27,7 @@ def min_heap_enqueue(A: ControlledPriorityQueue, x: KeyObject, n: int) -> None:
     A.priority += 1
 
 
-def min_heap_dequeue(A: ControlledPriorityQueue) -> KeyObject:
+def min_heap_dequeue(A: ControlledPriorityQueue[T]) -> KeyObject[T]:
     """Deletes an element from a first-in, first-out queue implemented with a priority queue.
 
     Implements:
@@ -41,7 +42,7 @@ def min_heap_dequeue(A: ControlledPriorityQueue) -> KeyObject:
     return min_heap_extract_min(A)
 
 
-def max_heap_push(A: ControlledPriorityQueue, x: KeyObject, n: int) -> None:
+def max_heap_push(A: ControlledPriorityQueue[T], x: KeyObject[T], n: int) -> None:
     """Inserts an element into a stack implemented with a priority queue.
 
     Implements:
@@ -57,7 +58,7 @@ def max_heap_push(A: ControlledPriorityQueue, x: KeyObject, n: int) -> None:
     A.priority += 1
 
 
-def max_heap_pop(A: ControlledPriorityQueue) -> KeyObject:
+def max_heap_pop(A: ControlledPriorityQueue[T]) -> KeyObject[T]:
     """Deletes an element from a stack implemented with a priority queue.
 
     Implements: