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

题目描述

这是 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]$ 范围内
• $-10^5 <= Node.val <= 10^5$

---

中序遍历 + 归并排序

利用 BST 中序遍历的有序性质，我们可以先对两棵树进行中序遍历，从而将树的结构转换为线性结构。

将两个有序序列合并成一个有序序列则是利用了经典的「归并排序」。

Java 代码：

class Solution {
    int INF = 0x3f3f3f3f;
    public List<Integer> getAllElements(TreeNode root1, TreeNode root2) {
        List<Integer> ans = new ArrayList<>();
        List<Integer> l1 = new ArrayList<>(), l2 = new ArrayList<>();
        dfs(root1, l1); dfs(root2, l2);
        int n = l1.size(), m = l2.size(), i = 0, j = 0;
        while (i < n || j < m) {
            int a = i < n ? l1.get(i) : INF, b = j < m ? l2.get(j) : INF;
            if (a <= b) {
                ans.add(a); i++;
            } else {
                ans.add(b); j++;
            }
        }
        return ans;
    }
    void dfs(TreeNode root, List<Integer> list) {
        if (root == null) return ;
        dfs(root.left, list);
        list.add(root.val);
        dfs(root.right, list);
    }
}

C++ 代码：

class Solution {
public:
    vector<int> getAllElements(TreeNode* root1, TreeNode* root2) {
        int INF = 0x3f3f3f3f;
        vector<int> ans;
        vector<int> l1, l2;
        dfs(root1, l1); dfs(root2, l2);
        int n = l1.size(), m = l2.size(), i = 0, j = 0;
        while (i < n || j < m) {
            int a = (i < n) ? l1[i] : INF, b = (j < m) ? l2[j] : INF;
            if (a <= b) {
                ans.push_back(a); i++;
            } else {
                ans.push_back(b); j++;
            }
        }
        return ans;
    }
    void dfs(TreeNode* root, vector<int>& list) {
        if (!root) return;
        dfs(root->left, list);
        list.push_back(root->val);
        dfs(root->right, list);
    }
};

Python 代码：

class Solution:
    def getAllElements(self, root1: Optional[TreeNode], root2: Optional[TreeNode]) -> List[int]:
        INF = 0x3f3f3f3f
        ans = []
        l1, l2 = [], []
        self.dfs(root1, l1)
        self.dfs(root2, l2)
        n, m, i, j = len(l1), len(l2), 0, 0
        while i < n or j < m:
            a, b = l1[i] if i < n else INF, l2[j] if j < m else INF
            if a <= b:
                ans.append(a)
                i += 1
            else:
                ans.append(b)
                j += 1
        return ans
        
    def dfs(self, root: TreeNode, arr: list[int]) -> None:
        if not root:
            return
        self.dfs(root.left, arr)
        arr.append(root.val)
        self.dfs(root.right, arr)

TypeScript 代码：

function dfs(root: TreeNode | null, list: number[]): void {
    if (root === null) return;
    dfs(root.left, list);
    list.push(root.val);
    dfs(root.right, list);
}
function getAllElements(root1: TreeNode | null, root2: TreeNode | null): number[] {
    const INF = 0x3f3f3f3f;
    let ans = [];
    let l1 = [], l2 = [];
    dfs(root1, l1);
    dfs(root2, l2);
    let n = l1.length, m = l2.length, i = 0, j = 0;
    while (i < n || j < m) {
        let a = i < n ? l1[i] : INF;
        let b = j < m ? l2[j] : INF;
        if (a <= b) {
            ans.push(a); i++;
        } else {
            ans.push(b); j++;
        }
    }
    return ans;
};

• 时间复杂度：令 $n$ 和 $m$ 分别为两棵树的节点数量，跑中序遍历的复杂度为 $O(n + m)$，构建答案复杂度为 $O(\max(m, n))$。整体复杂度为 $O(n + m)$
• 空间复杂度：$O(n + m)$

---

最后

这是我们「刷穿 LeetCode」系列文章的第 No.1305 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。

在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。

为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：https://github.com/SharingSource/LogicStack-LeetCode 。

在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。