[路飞]两颗二叉搜索树中的所有元素

102 阅读2分钟

记录 1 道算法题

两颗二叉搜索树中的所有元素

leetcode-cn.com/problems/al…


要求:将两颗二叉搜索树合并成一个升序的数组。

这道题如果利用了二叉搜索树的特点,解题是特别简单的。首先二叉搜索树由于左边的子节点小于根节点,右边的子节点大于根节点,所有中序遍历的时候可以得到一个升序的数组。那么对两个升序的数组进行合并,我们可以使用归并。

首先中序遍历,使用递归的话代码很短也很简单,这里使用循环来实现中序遍历。

    function dfs(root) {
        let result = []
        let stack = []
        // 通过根节点和栈来判断当前是处理当前应该递归深入还是往上返回
        while(root || stack.length) {
            // 我们假设按照递归的思路,首先一进来先进入到左子节点的深处
            while(root) {
                stack.push(root)
                root = root.left
            }
            // 返回的是最后的左子节点
            const node = stack.pop()
            result.push(node.val)
            
            // 到了最尽头的左子节点,开始向上返回,
            // 这时候有两种情况,我们回到根节点,或者进入右子节点
            // 因为当前递归的左节点尽头可能有一个右节点。
            if (node.right) {
                // 将右节点视为根节点,继续循环的步骤,进入根节点的左子节点深处
                root = node.right
            } else {
                // 设置为 null 之后就不会在下一轮寻找左子节点的深处,而是向上返回
                root = null
            }
        }
    }
    
    function getAllElements(root1, root2) {
        const arr1 = dfs(root1)
        const arr2 = dfs(root2)
        
        // 这里拿到了两个有序数组,进行归并
        let res = []
        let a = 0
        let b = 0
        // 比较最小的一位,然后指针后移
        while(a < arr1.length && b < arr2.length) {
            if (arr1[a] < arr2[b]) {
                res.push(arr1[a])
                a++
            } else {
                res.push(arr2[b])
                b++
            }
        }
        // 数组长短不一,可能有一个没有处理完,用 for 只是运行快而已
        if (a < arr1.length) {
            for(let i = a; i < arr1.length; i++) {
                res.push(arr1[i])
            }
        } else if (b < arr2.length) {
            for(let i = b; i < arr2.length; i++) {
                res.push(arr2[i])
            }
        }
        
        return res
    }