LeetCode173: 二叉搜索树的最小元素迭代器

问题描述

给定一个二叉搜索树，实现两个函数，一个是next()用来依次返回树中最小的元素，另一个是hasNext()用来判断是否还有下一个元素。类的框架如下所示。

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class BSTIterator {
    
    public BSTIterator(TreeNode root) {
        
    }
   
    /** @return the next smallest number */
    public int next() {
        
    }
    
    /** @return whether we have a next smallest number */
    public boolean hasNext() {
        
    }
}

/**
 * Your BSTIterator object will be instantiated and called as such:
 * BSTIterator obj = new BSTIterator(root);
 * int param_1 = obj.next();
 * boolean param_2 = obj.hasNext();
 */

问题分析

解法一：

将二叉搜索树按照中序遍历存进一个队列中，依次出队就能实现next()功能。不过，本文的重点是想介绍方法二。

解法二：

依然在树上进行查找，在遍历的过程中，修改叶子节点的右指针，指向该节点在中序遍历中的直接后继结点。具体过程如下图所示：

class BSTIterator {
    TreeNode node;
    public BSTIterator(TreeNode root) {
        this.node = root;
    }

    /** @return the next smallest number */
    public int next() {
        int ans;
        while(node.left != null) {
            TreeNode tmp = node.left;
            while(tmp.right != null && tmp.right != node) {
                tmp = tmp.right;
            }
            if(tmp.right == node) {
                tmp.right = null;
                break;
            } else {
                tmp.right = node;
                node = node.left;
            }
        }
        // 最左边就是答案
        ans = node.val;
        // 下次从答案的右子树开始找
        node = node.right;
        return ans;
    }

    /** @return whether we have a next smallest number */
    public boolean hasNext() {
        return node != null;
    }
}