Add comprehensive doctests to disjoint_set.py

data_structures/disjoint_set/disjoint_set.py

@@ -3,17 +3,28 @@
 Reference: https://en.wikipedia.org/wiki/Disjoint-set_data_structure
 """
 
+from dataclasses import dataclass
 
+
+@dataclass
 class Node:
-    def __init__(self, data: int) -> None:
-        self.data = data
-        self.rank: int
-        self.parent: Node
+    data: int
+    rank: int = 0
+    parent: Node | None = None
 
 
 def make_set(x: Node) -> None:
     """
     Make x as a set.
+
+    >>> node = Node(1)
+    >>> make_set(node)
+    >>> node.parent == node
+    True
+    >>> node.rank
+    0
+    >>> node.data
+    1
     """
     # rank is the distance from x to its' parent
     # root's rank is 0
@@ -26,6 +37,22 @@ def union_set(x: Node, y: Node) -> None:
     Union of two sets.
     set with bigger rank should be parent, so that the
     disjoint set tree will be more flat.
+
+    >>> node1 = Node(1)
+    >>> node2 = Node(2)
+    >>> make_set(node1)
+    >>> make_set(node2)
+    >>> union_set(node1, node2)
+    >>> find_set(node1) == find_set(node2)
+    True
+    >>> # Test union of already connected nodes
+    >>> node3 = Node(3)
+    >>> make_set(node3)
+    >>> union_set(node1, node3)
+    >>> find_set(node1) == find_set(node3)
+    True
+    >>> find_set(node2) == find_set(node3)
+    True
     """
     x, y = find_set(x), find_set(y)
     if x == y:
@@ -42,6 +69,24 @@ def union_set(x: Node, y: Node) -> None:
 def find_set(x: Node) -> Node:
     """
     Return the parent of x
+
+    >>> node = Node(1)
+    >>> make_set(node)
+    >>> find_set(node) == node
+    True
+    >>> node1 = Node(1)
+    >>> node2 = Node(2)
+    >>> make_set(node1)
+    >>> make_set(node2)
+    >>> union_set(node1, node2)
+    >>> find_set(node1) == find_set(node2)
+    True
+    >>> # Test path compression
+    >>> node3 = Node(3)
+    >>> make_set(node3)
+    >>> union_set(node1, node3)
+    >>> find_set(node1) == find_set(node3)
+    True
     """
     if x != x.parent:
         x.parent = find_set(x.parent)
@@ -51,6 +96,18 @@ def find_set(x: Node) -> Node:
 def find_python_set(node: Node) -> set:
     """
     Return a Python Standard Library set that contains i.
+
+    >>> node = Node(1)
+    >>> find_python_set(node)
+    {0, 1, 2}
+    >>> node = Node(4)
+    >>> find_python_set(node)
+    {3, 4, 5}
+    >>> node = Node(6)
+    >>> find_python_set(node)
+    Traceback (most recent call last):
+       ...
+    ValueError: 6 is not in ({0, 1, 2}, {3, 4, 5})
     """
     sets = ({0, 1, 2}, {3, 4, 5})
     for s in sets:
@@ -62,6 +119,9 @@ def find_python_set(node: Node) -> set:
 
 def test_disjoint_set() -> None:
     """
+    Test the disjoint set operations with a comprehensive example.
+    Creates two disjoint sets: {0, 1, 2} and {3, 4, 5}
+
     >>> test_disjoint_set()
     """
     vertex = [Node(i) for i in range(6)]