[mypy] Added type annotations to disjoint_set.py

data_structures/disjoint_set/disjoint_set.py

@@ -1,27 +1,29 @@
 """
-    disjoint set
+    Disjoint set.
     Reference: https://en.wikipedia.org/wiki/Disjoint-set_data_structure
 """
 
 
 class Node:
-    def __init__(self, data):
+    def __init__(self, data: int) -> None:
         self.data = data
+        self.rank: int
+        self.parent: Node
 
 
-def make_set(x):
+def make_set(x: Node) -> None:
     """
-    make x as a set.
+    Make x as a set.
     """
     # rank is the distance from x to its' parent
     # root's rank is 0
     x.rank = 0
     x.parent = x
 
 
-def union_set(x, y):
+def union_set(x: Node, y: Node) -> None:
     """
-    union two sets.
+    Union of two sets.
     set with bigger rank should be parent, so that the
     disjoint set tree will be more flat.
     """
@@ -37,9 +39,9 @@ def union_set(x, y):
             y.rank += 1
 
 
-def find_set(x):
+def find_set(x: Node) -> Node:
     """
-    return the parent of x
+    Return the parent of x
     """
     if x != x.parent:
         x.parent = find_set(x.parent)
@@ -57,7 +59,7 @@ def find_python_set(node: Node) -> set:
     raise ValueError(f"{node.data} is not in {sets}")
 
 
-def test_disjoint_set():
+def test_disjoint_set() -> None:
     """
     >>> test_disjoint_set()
     """