230 Kth Smallest Element in a BST

230. Kth Smallest Element in a BST

1. Question

Given a binary search tree, write a functionkthSmallestto find thekth 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 kth Smallest routine?

Hint:

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

2. Implementation

(1) Inorder traversal

思路:直接中序遍历找出第K小的node

class Solution {
    public int kthSmallest(TreeNode root, int k) {
        int count = 1;
        Stack<TreeNode> stack = new Stack<>();
        TreeNode curNode = root;

        while (curNode != null || !stack.isEmpty()) {
            if (curNode != null) {
                stack.push(curNode);
                curNode = curNode.left;
            }
            else {
                curNode = stack.pop();
                if (count == k) {
                    return curNode.val;
                }
                ++count;
                curNode = curNode.right;
            }
        }
        return -1;
    }
}

(2) Binary Search

思路:利用bst的二分特性,计算一个node在树里有多少个nodes的值比它小,如果该值等于k + 1, 说明当前的node是第k小的node,如果该值小于k + 1, 则第k小的node会在当前node的左边,否则在右边

public int kthSmallest(TreeNode root, int k) {
        int total = count(root.left);

        if (k <= total) {
            return kthSmallest(root.left, k);
        }
        else if (k > total + 1) {
            return kthSmallest(root.right, k - total - 1);
        }
        else {
            return root.val;
        }
    }

    public int count(TreeNode node) {
        if (node == null) {
            return 0;
        }
        return 1 + count(node.left) + count(node.right);
    }

(3) Follow up

如果我们频繁的操作该树,并且频繁的调用kth函数,有什么优化方法使时间复杂度降低至O(h)?h是树的高度。根据提示,我们可以在TreeNode中加入一个rank成员,这个变量记录的是该节点的左子树中节点的个数,其实就是有多少个节点比该节点小。这样我们就可以用二叉树搜索的方法来解决这个问题了。这个添加rank的操作可以在建树的时候一起完成。

3. Time & Space Complexity

Inorder Traversal: 时间复杂度: O(n), 空间复杂度: O(h), h为树的高度

Binary Search: 时间复杂度: O(h), 空间复杂度:O(h), count()这个function是递归地call,call stack最多可达树的高度

Last updated