Added type-hints and docs

Author: lucasimi

src/tdamapper/utils/heap.py

@@ -1,22 +1,47 @@
-def _left(i):
+"""
+This module implements a max-heap data structure.
+It provides methods to add elements, pop the maximum element,
+and iterate over the elements in the heap.
+"""
+
+
+def _left(i: int) -> int:
     return 2 * i + 1
 
 
-def _right(i):
+def _right(i: int) -> int:
     return 2 * i + 2
 
 
-def _parent(i):
+def _parent(i: int) -> int:
     return max(0, (i - 1) // 2)
 
 
 class _HeapNode:
+    """
+    A node in the max-heap, storing a key and a value.
+
+    It provides methods to get the key and value, and to compare nodes.
+    The comparison is based on the key, allowing the heap to maintain
+    the max-heap property. It is important to note that the key must be
+    comparable with other keys in the heap for the heap to function correctly.
+
+    :param key: The key used for ordering in the heap.
+    :param value: The value associated with the key.
+    :type key: Any
+    :type value: Any
+    """
 
     def __init__(self, key, value):
         self._key = key
         self._value = value
 
     def get(self):
+        """
+        Returns the key and value of the node.
+        :return: A tuple containing the key and value.
+        :rtype: tuple
+        """
         return self._key, self._value
 
     def __lt__(self, other):
@@ -33,6 +58,12 @@ def __ge__(self, other):
 
 
 class MaxHeap:
+    """
+    A max-heap data structure that allows for efficient retrieval of the maximum element.
+
+    It supports adding elements, popping the maximum element, and iterating over the elements.
+    It is important to note that the keys used in the heap must be comparable with each other.
+    """
 
     def __init__(self):
         self._heap = []
@@ -50,11 +81,23 @@ def __len__(self):
         return len(self._heap)
 
     def top(self):
+        """
+        Returns the maximum element in the heap without removing it.
+
+        :return: A tuple containing the key and value of the maximum element, or (None, None) if the heap is empty.
+        :rtype: tuple
+        """
         if not self._heap:
             return (None, None)
         return self._heap[0].get()
 
     def pop(self):
+        """
+        Removes and returns the maximum element from the heap.
+
+        :return: A tuple containing the key and value of the maximum element, or None if the heap is empty.
+        :rtype: tuple
+        """
         if not self._heap:
             return
         max_val = self._heap[0]
@@ -64,6 +107,16 @@ def pop(self):
         return max_val.get()
 
     def add(self, key, val):
+        """
+        Adds a new element to the heap with the specified key and value.
+
+        :param key: The key used for ordering in the heap.
+        :param val: The value associated with the key.
+        :type key: Any
+        :type val: Any
+        :return: None
+        :rtype: None
+        """
         self._heap.append(_HeapNode(key, val))
         self._bubble_up()
 

src/tdamapper/utils/unionfind.py

@@ -1,10 +1,40 @@
-class UnionFind:
+"""
+This module implements a Union-Find (Disjoint Set Union) data structure.
+It provides methods to find the root of an element and to union two elements.
+It is useful for efficiently managing and merging disjoint sets.
+"""
 
-    def __init__(self, X):
-        self._parent = {x: x for x in X}
+from typing import Dict, Generic, Iterable, TypeVar
+
+T = TypeVar("T")
+
+
+class UnionFind(Generic[T]):
+    """
+    A Union-Find data structure for managing disjoint sets.
+    It supports efficient union and find operations, allowing for the
+    merging of sets and finding the root representative of an element.
+    It is particularly useful in algorithms that require dynamic connectivity,
+    such as Kruskal's algorithm for minimum spanning trees or in network connectivity problems.
+
+    :param X: An iterable containing the initial elements to be managed by the Union-Find structure.
+    :type X: iterable
+    """
+
+    def __init__(self, X: Iterable[T]):
+        self._parent: Dict[T, T] = {x: x for x in X}
         self._size = {x: 1 for x in X}
 
-    def find(self, x):
+    def find(self, x: T) -> T:
+        """
+        Finds the root representative of the set containing element x.
+        This method implements path compression to flatten the structure,
+        making future queries faster by ensuring that all elements point directly to the root.
+        :param x: The element whose root representative is to be found.
+        :type x: Any
+        :return: The root representative of the set containing x.
+        :rtype: Any
+        """
         root = x
         while root != self._parent[root]:
             root = self._parent[root]
@@ -15,7 +45,19 @@ def find(self, x):
             tmp = parent
         return root
 
-    def union(self, x, y):
+    def union(self, x: T, y: T) -> T:
+        """
+        Unites the sets containing elements x and y.
+        If the elements are already in the same set, no action is taken.
+        If they are in different sets, the smaller set is merged into the larger set,
+        ensuring that the root of the larger set becomes the new root representative.
+        :param x: The first element to be united.
+        :param y: The second element to be united.
+        :type x: Any
+        :type y: Any
+        :return: The root representative of the united set.
+        :rtype: Any
+        """
         x, y = self.find(x), self.find(y)
         if x != y:
             x_size, y_size = self._size[x], self._size[y]