Swift - LeetCode - 二叉搜索树中的众数

108 阅读6分钟

持续创作,加速成长!这是我参与「掘金日新计划 · 10 月更文挑战」的第21天,点击查看活动详情

题目

给你一个含重复值的二叉搜索树(BST)的根节点 root,找出并返回 BST 中的所有 众数(即,出现频率最高的元素)。

如果树中有不止一个众数,可以按 任意顺序 返回。

假定 BST 满足如下定义:

  • 结点左子树中所含节点的值 小于等于 当前节点的值
  • 结点右子树中所含节点的值 大于等于 当前节点的值
  • 左子树和右子树都是二叉搜索树

示例 1:

image.png

  • 输入: root = [1,null,2,2]
  • 输出: [2]

示例 2:

  • 输入: root = [0]
  • 输出: [0]

方法一:中序遍历

思路及解法

首先我们一定能想到一个最朴素的做法:因为这棵树的中序遍历是一个有序的序列,所以我们可以先获得这棵树的中序遍历,然后从扫描这个中序遍历序列,然后用一个哈希表来统计每个数字出现的个数,这样就可以找到出现次数最多的数字。但是这样做的空间复杂度显然不是 O(1)O(1) 的,原因是哈希表和保存中序遍历序列的空间代价都是 O(n)O(n)

首先,我们考虑在寻找出现次数最多的数时,不使用哈希表。 这个优化是基于二叉搜索树中序遍历的性质:一棵二叉搜索树的中序遍历序列是一个非递减的有序序列。例如:

      1
    /   \
   0     2
  / \    /
-1   0  2

这样一颗二叉搜索树的中序遍历序列是 {1,0,0,1,2,2}\{ -1, 0, 0, 1, 2, 2 \}。我们可以发现重复出现的数字一定是一个连续出现的,例如这里的 0022,它们都重复出现了,并且所有的 00 都集中在一个连续的段内,所有的 22 也集中在一个连续的段内。我们可以顺序扫描中序遍历序列,用 base\textit{base} 记录当前的数字,用 count\textit{count} 记录当前数字重复的次数,用 maxCount\textit{maxCount} 来维护已经扫描过的数当中出现最多的那个数字的出现次数,用 answer\textit{answer} 数组记录出现的众数。每次扫描到一个新的元素:

  • 首先更新 base\textit{base}count\textit{count}:
    • 如果该元素和 base\textit{base} 相等,那么 count\textit{count} 自增 1;
    • 否则将 base\textit{base} 更新为当前数字,count\textit{count} 复位为 1。
  • 然后更新 maxCount\textit{maxCount}
    • 如果 count=maxCount\textit{count} = maxCount,那么说明当前的这个数字(base\textit{base})出现的次数等于当前众数出现的次数,将 base\textit{base} 加入 answer\textit{answer} 数组;
    • 如果 count>maxCount\textit{count} > maxCount,那么说明当前的这个数字(base\textit{base})出现的次数大于当前众数出现的次数,因此,我们需要将 maxCount\textit{maxCount} 更新为 count\textit{count},清空 answer\textit{answer} 数组后将 base\textit{base} 加入 answer\textit{answer} 数组。

我们可以把这个过程写成一个 update\text{update} 函数。这样我们在寻找出现次数最多的数字的时候就可以省去一个哈希表带来的空间消耗。

然后,我们考虑不存储这个中序遍历序列。 如果我们在递归进行中序遍历的过程中,访问当了某个点的时候直接使用上面的 update\text{update} 函数,就可以省去中序遍历序列的空间,代码如下。

代码

class Solution {
    var answer: [Int] = []
    var base: Int = 0
    var count: Int = 0
    var maxCount: Int = 0
    
    func findMode(_ root: TreeNode?) -> [Int] {
        dfs(root)
        return answer
    }
    
    func dfs(_ o: TreeNode?) {
        if nil == o {
            return
        }
        dfs(o?.left)
        update(o!.val)
        dfs(o?.right)
    }
    
    func update(_ x: Int) {
        if x == base {
            count += 1
        } else {
            count = 1
            base = x
        }
        if count == maxCount {
            answer.append(base)
        }
        if count > maxCount {
            maxCount = count
            answer = []
            answer.append(base)
        }
    }
}

复杂度分析

  • 时间复杂度:O(n)O(n)。即遍历这棵树的复杂度。

  • 空间复杂度:O(n)O(n)。即递归的栈空间的空间代价。

方法二:Morris 中序遍历

思路及解法

接着上面的思路,我们用 Morris 中序遍历的方法把中序遍历的空间复杂度优化到 O(1)O(1)。 我们在中序遍历的时候,一定先遍历左子树,然后遍历当前节点,最后遍历右子树。在常规方法中,我们用递归回溯或者是栈来保证遍历完左子树可以再回到当前节点,但这需要我们付出额外的空间代价。我们需要用一种巧妙地方法可以在 O(1)O(1) 的空间下,遍历完左子树可以再回到当前节点。我们希望当前的节点在遍历完当前点的前驱之后被遍历,我们可以考虑修改它的前驱节点的 right\textit{right} 指针。当前节点的前驱节点的right \textit{right} 指针可能本来就指向当前节点(前驱是当前节点的父节点),也可能是当前节点左子树最右下的节点。如果是后者,我们希望遍历完这个前驱节点之后再回到当前节点,可以将它的 right\textit{right} 指针指向当前节点。

Morris 中序遍历的一个重要步骤就是寻找当前节点的前驱节点,并且 Morris 中序遍历寻找下一个点始终是通过转移到 right\textit{right} 指针指向的位置来完成的。

  • 如果当前节点没有左子树,则遍历这个点,然后跳转到当前节点的右子树。
  • 如果当前节点有左子树,那么它的前驱节点一定在左子树上,我们可以在左子树上一直向右行走,找到当前点的前驱节点。
    • 如果前驱节点没有右子树,就将前驱节点的 right\textit{right} 指针指向当前节点。这一步是为了在遍历完前驱节点后能找到前驱节点的后继,也就是当前节点。
    • 如果前驱节点的右子树为当前节点,说明前驱节点已经被遍历过并被修改了 right\textit{right} 指针,这个时候我们重新将前驱的右孩子设置为空,遍历当前的点,然后跳转到当前节点的右子树。

代码

extension TreeNode: Equatable {
    public static func == (lhs: TreeNode, rhs: TreeNode) -> Bool {
        return lhs === rhs
    }
}

class Solution {
    var answer: [Int] = []
    var base: Int = 0
    var count: Int = 0
    var maxCount: Int = 0
    
    func update(_ x: Int) {
        if x == base {
            count += 1
        } else {
            count = 1
            base = x
        }
        if count == maxCount {
            answer.append(base)
        }
        if count > maxCount {
            maxCount = count
            answer = [base]
        }
    }
    
    func findMode(_ root: TreeNode?) -> [Int] {
        var cur: TreeNode? = root
        var pre: TreeNode? = nil
        while nil != cur {
            if nil == cur?.left {
                update(cur!.val)
                cur = cur?.right
                continue
            }
            pre = cur?.left
            while nil != pre?.right && pre?.right != cur {
                pre = pre?.right
            }
            if nil == pre?.right {
                pre?.right = cur
                cur = cur?.left
            } else {
                pre?.right = nil
                update(cur!.val)
                cur = cur?.right
            }
        }
        return answer
    }
}

复杂度分析

  • 时间复杂度:O(n)O(n)。每个点被访问的次数不会超过两次,故这里的时间复杂度是 O(n)O(n)

  • 空间复杂度:O(1)O(1)。使用临时空间的大小和输入规模无关。