题目描述

给定一个二叉搜索树，编写一个函数 kthSmallest 来查找其中第 k 个最小的元素。

说明：
你可以假设 k 总是有效的，1 ≤ k ≤ 二叉搜索树元素个数。

示例 1:

输入: root = [3,1,4,null,2], k = 1
   3
  / \
 1   4
  \
   2
输出: 1
示例 2:

输入: root = [5,3,6,2,4,null,null,1], k = 3
       5
      / \
     3   6
    / \
   2   4
  /
 1
输出: 3

解法

中序遍历

因为是二叉搜索树，树的结构是有序的。通过中序遍历可以得到一个有序数组，求第k个元素即可。

public int kthSmallest(TreeNode root, int k) {
        List<Integer> result = new ArrayList<>();
        inOrder(root, result);
        return result.get(k-1);
    }

    public void inOrder(TreeNode node, List<Integer> result) {
        if (node == null) {
            return;
        }
        inOrder(node.left, result);
        result.add(node.val);
        inOrder(node.right, result);
    }

运行结果：

中序遍历-剪枝

如果已经遍历到第k大的元素，就没必要遍历完整个树

int count = 0;
    int result = 0;
    public int kthSmallest(TreeNode root, int k) {
       inOrder(root, k);
       return result;
    }

    public void inOrder(TreeNode node, int k) {
        if (node == null) {
            return;
        }
        if (node.left != null) {
            inOrder(node.left, k);
        }
        count ++;
        if (count == k) {
            result = node.val;
        }
        if(count > k) {
            return;
        }

        if (node.right != null) {
            inOrder(node.right, k);
        }
    }

迭代方式

利用栈来解决。先将左子树压入栈，再开始回弹，遍历右子树，判断是否是第k个

public int kthSmallest(TreeNode root, int k) {

    LinkedList<TreeNode> stack = new LinkedList<TreeNode>();

    while (true) {
      while (root != null) {
        stack.add(root);
        root = root.left;
      }
      root = stack.removeLast();
      if (--k == 0) return root.val;
      root = root.right;
    }
  }