From abaa0d754b8dc24abacac2a4d7ecade2d3ddacb6 Mon Sep 17 00:00:00 2001
From: scfenton6 <91698851+scfenton6@users.noreply.github.com>
Date: Mon, 11 Oct 2021 18:34:30 +0200
Subject: [PATCH] Add type annotations (#4814)

---
 data_structures/disjoint_set/disjoint_set.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/data_structures/disjoint_set/disjoint_set.py b/data_structures/disjoint_set/disjoint_set.py
index bf5ab415d..f8500bf2c 100644
--- a/data_structures/disjoint_set/disjoint_set.py
+++ b/data_structures/disjoint_set/disjoint_set.py
@@ -1,17 +1,19 @@
 """
-    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
@@ -19,9 +21,9 @@ def make_set(x):
     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()
     """