solved 'Kth Smallest Element in a BST' on LeetCode https://leetcode.com/problems/kth-smallest-element-in-a-bst/

kth_bst_bad.py

# Memoize the size and iteratively get down.
# Bad because getting size requires O(n) time.
class Solution(object):
    def kthSmallest(self, root, k):
        """
        :type root: TreeNode
        :type k: int
        :rtype: int
        """
        self.sz_dict = {None: 0}
        
        while root:
            sz_left = self.get_size(root.left)
            if sz_left == k - 1:
                return root.val
            elif sz_left >= k:
                root = root.left
            else:
                root = root.right
                k -= 1 + sz_left
        
        # Should never reach here.
        return -1
        
    def get_size(self, root):
        if not root:
            return 0
            
        sz = self.sz_dict.get(root, -1)
        if sz == -1:
            sz = self.get_size(root.left) + self.get_size(root.right) + 1
            self.sz_dict[root] = sz
            
        return sz

kth_bst_good.py

# In order traversal using stack. K-th iteration returns
# k-th smallest.
class Solution(object):
    def kthSmallest(self, root, k):
        """
        :type root: TreeNode
        :type k: int
        :rtype: int
        """
        # Use stack for in-order traversal.
        st, curr = [], root
        # Since k is always valid, so the
        # loop is guaranteed to end.
        while True:
            if curr:
                st.append(curr)
                curr = curr.left
            else:
                curr = st.pop()
                k -= 1
                # Matched k-th smallest.
                if k == 0:
                    return curr.val
                
                curr = curr.right