230.Kth Smallest Element in a BST

Tags: [tree], [binary_tree], [binary_search_tree], [BST], [inorder], [stack]

Com: {g}

Hard: [###]

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

Given a binary search tree, write a function kthSmallestto 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?

Solution: Inorder Traversal, Stack

# 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
        """
        if not root or not k:
            return None

        counter = 0
        stack = []
        curr = root
        while curr or stack:
            if curr:
                stack.append(curr)
                curr = curr.left
            else:
                curr = stack.pop()
                counter += 1
                if counter == k:
                    return curr.val

                curr = curr.right

        return None

Note:

• Time complexity = O(n), n is the number of nodes in the given BST.

Follow Up:

# 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):
    class CountTreeNode(object):
        def __init__(self, val):
            self.val = val
            self.count = 1
            self.left = None
            self.right = None


    def kthSmallest(self, root, k):
        """
        :type root: TreeNode
        :type k: int
        :rtype: int
        """
        if not root or not k:
            return None

        count_root = self.build_count_tree(root)
        return self.helper(count_root, k)

    def build_count_tree(self, root):
        if not root:
            return None

        count_root = self.CountTreeNode(root.val)
        count_root.left = self.build_count_tree(root.left)
        count_root.right = self.build_count_tree(root.right)

        if count_root.left:
            count_root.count += count_root.left.count
        if count_root.right:
            count_root.count += count_root.right.count

        return count_root

    def helper(self, count_root, k):
        # base case
        if not count_root or k < 0 or k > count_root.count:
            return -1

        if count_root.left:
            if count_root.left.count >= k:
                return self.helper(count_root.left, k)
            if count_root.left.count == k - 1:
                return count_root.val
            return self.helper(count_root.right, k - 1 - count_root.left.count)
        else:
            if k == 1:
                return count_root.val
            return self.helper(count_root.right, k - 1)

Revelation:

• If we can modify the TreeNode structure, we can let each node contains a count, which contains the number of nodes in the current sub tree.
• To build this count tree need T = O(n), n is the number of nodes in the given tree. But after each insert / delete operations, it only need T = O(lgn) to get the kth smallest element.