【Lintcode】902. Kth Smallest Element in a BST

题目地址：

https://www.lintcode.com/problem/kth-smallest-element-in-a-bst/description

给定一棵BST，求其第 $k$小的元素。

和Quick Select算法类似，可以先确定root排在第几小，根据情况再判断答案在左子树还是右子树，再将指针向左子树或者右子树移动。为了避免重复计算，可以用一个哈希表记录每个节点有多少子节点。确定root排第几小的方法就是直接去算左子树有多少个节点。这里可以用递归的方法算，边算边存进哈希表。代码如下：

import java.util.HashMap;
import java.util.Map;

public class Solution {
    /**
     * @param root: the given BST
     * @param k: the given k
     * @return: the kth smallest element in BST
     */
    public int kthSmallest(TreeNode root, int k) {
        // write your code here
        // 记录每个节点有多少后代节点
        Map<TreeNode, Integer> map = new HashMap<>();
        int x = 0;
        // 如果左子树的节点个数等于k - 1，那就说明root即为所求，这时候就跳出循环
        // 如果左子树节点个数大于k - 1，说明第k小的节点在左子树，且正好是左子树的第k小的数，所以直接将指针左移；
        // 如果左子树节点个数大于k - 1，说明第k小的节点在右子树，且为右子树的第k - (x + 1)小的数，所以将指针右移并更新k
        while ((x = count(root.left, map)) != k - 1) {
            if (x > k - 1) {
                root = root.left;
            } else {
                root = root.right;
                k -= x + 1;
            }
        }
        
        return root.val;
    }
    
    private int count(TreeNode root, Map<TreeNode, Integer> map) {
        if (root == null) {
            return 0;
        }
        
        if (map.containsKey(root)) {
            return map.get(root);
        }
        
        int left = count(root.left, map);
        int right = count(root.right, map);
        map.put(root, left + right + 1);
        
        return left + right + 1;
    }
}

class TreeNode {
    int val;
    TreeNode left, right;
    TreeNode(int x) {
        val = x;
    }
}

时空复杂度 $O(n)$。