Medium

题目描述

给你两棵二叉搜索树 root1root2,请你返回一个列表,其中包含了两棵树中的所有整数并按升序排序。

示例 1:

输入:root1 = [2,1,4], root2 = [1,0,3]
输出:[0,1,1,2,3,4]

示例 2:

输入:root1 = [1,null,8], root2 = [8,1]
输出:[1,1,8,8]

提示:

  • 每棵树的节点数在 [0, 5000] 范围内
  • -10^5 <= Node.val <= 10^5

解题思路

这道题有两种主要思路:

思路1:分别遍历后合并排序 最直观的方法是分别对两棵BST进行中序遍历得到两个有序数组,然后将两个有序数组合并。由于BST的中序遍历天然有序,我们可以直接使用归并排序的思想来合并两个有序数组。

思路2:同时遍历两棵树 更优雅的方法是使用栈同时遍历两棵树。维护两个栈分别存储两棵树的遍历路径,每次比较栈顶元素的值,取较小者加入结果,然后继续遍历该树。

推荐解法:思路2 虽然两种方法的时间复杂度相同,但思路2在空间使用上更加灵活,不需要额外存储完整的中序遍历结果,而且代码更简洁。我们使用栈来模拟中序遍历的过程,通过比较两个栈顶元素来决定下一个要添加到结果中的元素。

具体实现时,我们先将两棵树最左侧的所有节点压入各自的栈中,然后在循环中比较栈顶元素,取较小者,并将其右子树的最左侧路径压入栈中。

代码实现

class Solution {
public:
    vector<int> getAllElements(TreeNode* root1, TreeNode* root2) {
        vector<int> result;
        stack<TreeNode*> s1, s2;
        
        while (root1 || root2 || !s1.empty() || !s2.empty()) {
            // 将左侧路径压入栈
            while (root1) {
                s1.push(root1);
                root1 = root1->left;
            }
            while (root2) {
                s2.push(root2);
                root2 = root2->left;
            }
            
            // 比较栈顶元素,选择较小的
            if (s2.empty() || (!s1.empty() && s1.top()->val <= s2.top()->val)) {
                root1 = s1.top();
                s1.pop();
                result.push_back(root1->val);
                root1 = root1->right;
            } else {
                root2 = s2.top();
                s2.pop();
                result.push_back(root2->val);
                root2 = root2->right;
            }
        }
        
        return result;
    }
};
class Solution:
    def getAllElements(self, root1: Optional[TreeNode], root2: Optional[TreeNode]) -> List[int]:
        result = []
        s1, s2 = [], []
        
        while root1 or root2 or s1 or s2:
            # 将左侧路径压入栈
            while root1:
                s1.append(root1)
                root1 = root1.left
            while root2:
                s2.append(root2)
                root2 = root2.left
            
            # 比较栈顶元素,选择较小的
            if not s2 or (s1 and s1[-1].val <= s2[-1].val):
                root1 = s1.pop()
                result.append(root1.val)
                root1 = root1.right
            else:
                root2 = s2.pop()
                result.append(root2.val)
                root2 = root2.right
        
        return result
public class Solution {
    public IList<int> GetAllElements(TreeNode root1, TreeNode root2) {
        var result = new List<int>();
        var s1 = new Stack<TreeNode>();
        var s2 = new Stack<TreeNode>();
        
        while (root1 != null || root2 != null || s1.Count > 0 || s2.Count > 0) {
            // 将左侧路径压入栈
            while (root1 != null) {
                s1.Push(root1);
                root1 = root1.left;
            }
            while (root2 != null) {
                s2.Push(root2);
                root2 = root2.left;
            }
            
            // 比较栈顶元素,选择较小的
            if (s2.Count == 0 || (s1.Count > 0 && s1.Peek().val <= s2.Peek().val)) {
                root1 = s1.Pop();
                result.Add(root1.val);
                root1 = root1.right;
            } else {
                root2 = s2.Pop();
                result.Add(root2.val);
                root2 = root2.right;
            }
        }
        
        return result;
    }
}
var getAllElements = function(root1, root2) {
    const result = [];
    const s1 = [], s2 = [];
    
    while (root1 || root2 || s1.length || s2.length) {
        // 将左侧路径压入栈
        while (root1) {
            s1.push(root1);
            root1 = root1.left;
        }
        while (root2) {
            s2.push(root2);
            root2 = root2.left;
        }
        
        // 比较栈顶元素,选择较小的
        if (!s2.length || (s1.length && s1[s1.length - 1].val <= s2[s2.length - 1].val)) {
            root1 = s1.pop();
            result.push(root1.val);
            root1 = root1.right;
        } else {
            root2 = s2.pop();
            result.push(root2.val);
            root2 = root2.right;
        }
    }
    
    return result;
};

复杂度分析

复杂度类型复杂度说明
时间复杂度O(m + n)m 和 n 分别是两棵树的节点数,需要遍历所有节点
空间复杂度O(h1 + h2)h1 和 h2 分别是两棵树的高度,栈的最大深度