LeetCode——1305. 两棵二叉搜索树中的所有元素

题目描述

给你 root1 和 root2 这两棵二叉搜索树。请你返回一个列表，其中包含 两棵树 中的所有整数并按 升序 排序。.

示例 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] 范围内
• -105 <= Node.val <= 105

答案

我的答案

class Solution {
    List<Integer> list = new ArrayList();
    public List<Integer> getAllElements(TreeNode root1, TreeNode root2) {
        dfs(root1);
        dfs(root2);
        Collections.sort(list);
        return list;
    }

    public void dfs(TreeNode root){
        if (root==null){
            return;
        }
        list.add(root.val);
        dfs(root.left);
        dfs(root.right);
    }
}

官方答案

中序遍历 + 归并

回顾二叉搜索树的定义：

• 当前节点的左子树中的数均小于当前节点的数；
• 当前节点的右子树中的数均大于当前节点的数；
• 所有左子树和右子树自身也是二叉搜索树。

根据上述定义，我们可以用中序遍历访问二叉搜索树，即按照访问左子树——根节点——右子树的方式遍历这棵树，而在访问左子树或者右子树的时候也按照同样的方式遍历，直到遍历完整棵树。遍历结束后，就得到了一个有序数组。

由于整个遍历过程天然具有递归的性质，我们可以直接用递归函数来模拟这一过程。具体描述见 94. 二叉树的中序遍历 的 官方题解。

中序遍历这两棵二叉搜索树，可以得到两个有序数组。然后可以使用双指针方法来合并这两个有序数组，这一方法将两个数组看作两个队列，每次从队列头部取出比较小的数字放到结果中（头部相同时可任取一个）。如下面的动画所示：

class Solution {
    public List<Integer> getAllElements(TreeNode root1, TreeNode root2) {
        List<Integer> nums1 = new ArrayList<Integer>();
        List<Integer> nums2 = new ArrayList<Integer>();
        inorder(root1, nums1);
        inorder(root2, nums2);

        List<Integer> merged = new ArrayList<Integer>();
        int p1 = 0, p2 = 0;
        while (true) {
            if (p1 == nums1.size()) {
                merged.addAll(nums2.subList(p2, nums2.size()));
                break;
            }
            if (p2 == nums2.size()) {
                merged.addAll(nums1.subList(p1, nums1.size()));
                break;
            }
            if (nums1.get(p1) < nums2.get(p2)) {
                merged.add(nums1.get(p1++));
            } else {
                merged.add(nums2.get(p2++));
            }
        }
        return merged;
    }

    public void inorder(TreeNode node, List<Integer> res) {
        if (node != null) {
            inorder(node.left, res);
            res.add(node.val);
            inorder(node.right, res);
        }
    }
}

复杂度分析

• 时间复杂度：O(n+m)，其中 n 和 m 分别为两棵二叉搜索树的节点个数。

• 空间复杂度：O(n+m)。存储数组以及递归时的栈空间均为 O(n+m)。