LeetCode-230. Kth Smallest Element in a BST

https://leetcode.com/problems/kth-smallest-element-in-a-bst/

Given a binary search tree, write a function kthSmallest to find the kth smallest element in it.

Note:
You may assume k is always valid, 1 ≤ k ≤ BST's total elements.

Follow up:
What if the BST is modified (insert/delete operations) often and you need to find the kth smallest frequently? How would you optimize the kthSmallest routine?

Hint:

1.Try to utilize the property of a BST.
2.What if you could modify the BST node's structure?
3.The optimal runtime complexity is O(height of BST).

Solution

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution(object):
    def kthSmallest(self, root, k):
        """
        :type root: TreeNode
        :type k: int
        :rtype: int
        
        count = self.countNode(root.left)
        if k <= count:
            return self.kthSmallest(root.left, k)
        elif k > count + 1:
            return self.kthSmallest(root.right, k-count-1)
        return root.val
        
    def countNode(self, node):
        if node == None:
            return 0
        return 1 + self.countNode(node.left) + self.countNode(node.right)
    """
        st = []
        while root is not None:
            st.append(root)
            root = root.left
        while k != 0:
            node = st.pop()
            print(node.val)
            k -= 1
            if k == 0:
                return node.val
            right = node.right
            while right:
                st.append(right)
                right = right.left
        return -1