change code to resolve errors

Author: sohamsshah

Diff for: data_structures/binary_tree/kth_smallest_value_in_a_binary_search_tree.py

@@ -12,18 +12,7 @@ def __init__(self, data: int) -> None:
         self.left: Optional[Node] = None
         self.right: Optional[Node] = None
 
-def display(tree: Optional[Node]) -> None:  # In Order traversal of the tree
-    """
-    >>> root = Node(1)
-    >>> root.left = Node(0)
-    >>> root.right = Node(2)
-    >>> display(root)
-    0
-    1
-    2
-    >>> display(root.right)
-    2
-    """
+def display(tree: Optional[Node]) -> None: 
     if tree:
         display(tree.left)
         print(tree.data)
@@ -33,7 +22,6 @@ def kthSmallest(root: Node, k: int) -> int:
     """
     In a BST, the Inorder Traversal returns ascending order of the data when traversed.
     Thus, store Inorder traversal in a list and return the k-1 th Index which is the Kth Smallest number in a BST.
-    
     """
     stack = []
     temp = []
@@ -47,17 +35,17 @@ def kthSmallest(root: Node, k: int) -> int:
             root = stack.pop()
             temp.append(root.data)
             root = root.right
-    return (temp[k-1])
+    return temp[k-1]
 
-def main() -> None:  # Main function for testing.
+def main() -> None: 
     tree = Node(5)
     tree.left = Node(3)
     tree.right = Node(6)
     tree.left.left = Node(2)
     tree.left.right = Node(4)
     tree.left.left.left = Node(1)
-    k = int(input("Enter the value of K:"))
-    print("The {}th Smallest Value in the BST is:".format(k), kthSmallest(tree, k))
+    k = 2  
+    print(f"The {k}th Smallest Value in the BST is:", kthSmallest(tree, k))
     print("Tree is: ")
     display(tree)