【Leetcode】1305. All Elements in Two Binary Search Trees

题目地址：

https://leetcode.com/problems/all-elements-in-two-binary-search-trees/

给定两棵二叉搜索树，返回由它们所有节点值构成的上升序列。

可以建一个二叉树中序遍历的迭代器，然后问题就转化为了合并两个有序链表。代码如下：

import java.util.ArrayList;
import java.util.Deque;
import java.util.LinkedList;
import java.util.List;

public class Solution {
     
    
    class Iterator {
     
        
        private Deque<TreeNode> stack;
        
        public Iterator(TreeNode root) {
     
            stack = new LinkedList<>();
            // 存一下左链
            while (root != null) {
     
                stack.push(root);
                root = root.left;
            }
        }
        
        public int peek() {
     
            return stack.peek().val;
        }
        
        public int next() {
     
        	// 每次pop节点的时候，就将右子树的左链进栈
            TreeNode cur = stack.pop(), right = cur.right;
            while (right != null) {
     
                stack.push(right);
                right = right.left;
            }
            
            return cur.val;
        }
        
        public boolean hasNext() {
     
            return !stack.isEmpty();
        }
    }
    
    public List<Integer> getAllElements(TreeNode root1, TreeNode root2) {
     
        List<Integer> res = new ArrayList<>();
        Iterator it1 = new Iterator(root1), it2 = new Iterator(root2);
        while (it1.hasNext() || it2.hasNext()) {
     
            if (!it1.hasNext()) {
     
                res.add(it2.next());
            } else if (!it2.hasNext()) {
     
                res.add(it1.next());
            } else {
     
                int v1 = it1.peek(), v2 = it2.peek();
                if (v1 < v2) {
     
                    res.add(it1.next());
                } else if (v1 > v2) {
     
                    res.add(it2.next());
                } else {
     
                    res.add(it1.next());
                    res.add(it2.next());
                }
            }
        }
        
        return res;
    }
}

class TreeNode {
     
    int val;
    TreeNode left, right;
    
    public TreeNode(int val) {
     
        this.val = val;
    }
}

均摊时间复杂度$O(n_1+n_2)$，空间$O(h_1+h_2)$。