Refactor: add type hints, doctests, and naming fixes for bot compliance

Author: cureprotocols

data_structures/disjoint_set/union_find.py

@@ -13,38 +13,64 @@
 
 
 class UnionFind:
-    def __init__(self, size):
+    def __init__(self, size: int) -> None:
+        """
+        Initializes a Union-Find data structure with `size` elements.
+
+        >>> uf = UnionFind(5)
+        >>> uf.find(0)
+        0
+        """
         self.parent = list(range(size))
         self.rank = [0] * size
 
-    def find(self, node):
+    def find(self, node: int) -> int:
+        """
+        Finds the representative/root of the set that `node` belongs to.
+
+        >>> uf = UnionFind(5)
+        >>> uf.find(3)
+        3
+        """
         if self.parent[node] != node:
             self.parent[node] = self.find(self.parent[node])  # Path compression
         return self.parent[node]
 
-    def union(self, x, y):
-        rootX = self.find(x)
-        rootY = self.find(y)
-
-        if rootX == rootY:
+    def union(self, a: int, b: int) -> bool:
+        """
+        Unites the sets that contain elements `a` and `b`.
+
+        >>> uf = UnionFind(5)
+        >>> uf.union(0, 1)
+        True
+        >>> uf.find(1) == uf.find(0)
+        True
+        >>> uf.union(0, 1)
+        False
+        """
+        root_a = self.find(a)
+        root_b = self.find(b)
+
+        if root_a == root_b:
             return False  # Already connected
 
         # Union by rank
-        if self.rank[rootX] < self.rank[rootY]:
-            self.parent[rootX] = rootY
-        elif self.rank[rootX] > self.rank[rootY]:
-            self.parent[rootY] = rootX
+        if self.rank[root_a] < self.rank[root_b]:
+            self.parent[root_a] = root_b
+        elif self.rank[root_a] > self.rank[root_b]:
+            self.parent[root_b] = root_a
         else:
-            self.parent[rootY] = rootX
-            self.rank[rootX] += 1
+            self.parent[root_b] = root_a
+            self.rank[root_a] += 1
 
         return True
 
 
-# Example usage
 if __name__ == "__main__":
-    uf = UnionFind(10)
+    import doctest
+    doctest.testmod()
 
+    uf = UnionFind(10)
     uf.union(1, 2)
     uf.union(2, 3)
     uf.union(4, 5)

searches/reservoir_sampling.py

@@ -2,48 +2,53 @@
 Reservoir Sampling Algorithm
 
 Use Case:
-Efficient for selecting k random items from a data stream of unknown size,
+Efficient for selecting `sample_size` random items from a data stream of unknown size,
 or when the entire dataset cannot fit into memory.
 
 Time Complexity:
 - O(n), where n is the total number of items
-- Space Complexity: O(k)
+- Space Complexity: O(sample_size)
 
 Author: Michael Alexander Montoya
 """
 
 import random
+from typing import Iterable
 
 
-def reservoir_sampling(stream, k):
+def reservoir_sampling(stream: Iterable[int], sample_size: int) -> list[int]:
     """
     Performs reservoir sampling on a stream of items.
 
     Args:
         stream: An iterable data stream.
-        k: Number of items to sample.
+        sample_size: Number of items to sample.
 
     Returns:
-        A list containing k randomly sampled items from the stream.
-    """
+        A list containing `sample_size` randomly sampled items from the stream.
 
+    >>> result = reservoir_sampling(range(1, 1001), 10)
+    >>> len(result) == 10
+    True
+    """
     reservoir = []
 
     for i, item in enumerate(stream):
-        if i < k:
+        if i < sample_size:
             reservoir.append(item)
         else:
             j = random.randint(0, i)
-            if j < k:
+            if j < sample_size:
                 reservoir[j] = item
 
     return reservoir
 
 
-# Example usage
 if __name__ == "__main__":
-    stream_data = range(1, 1001)  # Simulate a stream of numbers from 1 to 1000
-    sample_size = 10
+    import doctest
+
+    doctest.testmod()
 
-    sample = reservoir_sampling(stream_data, sample_size)
-    print(f"Random sample of {sample_size} items from stream: {sample}")
+    stream_data = range(1, 1001)
+    sample = reservoir_sampling(stream_data, 10)
+    print(f"Sampled items: {sample}")