LeetCode 501: Find Mode in Binary Search Tree

Problem Restatement

We are given the root of a binary search tree that may contain duplicate values.

We need to return all values that appear most frequently in the tree.

A value with the highest frequency is called a mode.

If more than one value has the same highest frequency, return all of them in any order. The official statement defines the input as a BST with duplicates and asks for all most frequently occurring elements.

Input and Output

The function shape is:

class Solution:
    def findMode(self, root: Optional[TreeNode]) -> List[int]:
        ...

Examples

Consider this tree:

    1
     \
      2
     /
    2

The values are:

1, 2, 2

The value 1 appears once.

The value 2 appears twice.

So the answer is:

[2]

Another example:

    2
   / \
  1   3

The values are:

1, 2, 3

Each value appears once.

So every value is a mode:

[1, 2, 3]

First Thought: Count Every Value

A simple solution is to traverse the whole tree and count each value with a hash map.

For example:

count[value] += 1

After traversal, we find the largest frequency and return every value whose frequency equals that maximum.

This works for any binary tree, not only a BST.

from collections import defaultdict
from typing import Optional, List

class Solution:
    def findMode(self, root: Optional[TreeNode]) -> List[int]:
        count = defaultdict(int)

        def dfs(node):
            if not node:
                return

            count[node.val] += 1
            dfs(node.left)
            dfs(node.right)

        dfs(root)

        max_count = max(count.values())
        return [value for value, freq in count.items() if freq == max_count]

Problem With This Approach

The hash map stores one entry for every distinct value.

If the tree has many distinct values, the extra space can be large.

Since the input is a binary search tree, we can do better by using its sorted traversal property.

Key Insight

An inorder traversal of a BST visits values in sorted order.

So duplicate values appear next to each other during traversal.

For example, this BST may produce:

1, 2, 2, 3, 3, 3, 4

Because equal values are consecutive, we only need to track the current value streak:

previous value
current count
maximum count
answer list

When the current value equals the previous value, increase the current count.

When the value changes, reset the current count to 1.

Algorithm

Use inorder traversal.

Maintain four pieces of state:

During inorder traversal:

1. Visit the left subtree.
2. Process the current node.
3. Visit the right subtree.

When processing a node:

1. If prev exists and prev.val == node.val, increment count.
2. Otherwise, reset count to 1.
3. If count > max_count, replace the answer with the current value.
4. If count == max_count, append the current value.
5. Set prev to the current node.

Correctness

Inorder traversal visits BST values in sorted order.

Therefore, all occurrences of the same value are visited as one contiguous group.

The algorithm counts the length of each such group using count.

When a group becomes larger than all previous groups, that value is the only known mode so far, so the answer list is replaced.

When a group has the same size as the largest known group, that value is also a mode, so it is appended.

After traversal finishes, every value group has been processed exactly once. The answer list contains exactly the values whose frequency equals the maximum frequency.

Complexity

The output list is not counted as auxiliary space.

For a balanced tree, h = O(log n).

For a skewed tree, h = O(n).

Implementation

from typing import Optional, List

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right

class Solution:
    def findMode(self, root: Optional[TreeNode]) -> List[int]:
        self.ans = []
        self.prev = None
        self.count = 0
        self.max_count = 0

        def update(node: TreeNode) -> None:
            if self.prev and self.prev.val == node.val:
                self.count += 1
            else:
                self.count = 1

            if self.count > self.max_count:
                self.max_count = self.count
                self.ans = [node.val]
            elif self.count == self.max_count:
                self.ans.append(node.val)

            self.prev = node

        def inorder(node: Optional[TreeNode]) -> None:
            if not node:
                return

            inorder(node.left)
            update(node)
            inorder(node.right)

        inorder(root)
        return self.ans

Code Explanation

The answer list starts empty:

self.ans = []

self.prev stores the previously visited node in inorder traversal:

self.prev = None

self.count stores how many times the current value has appeared consecutively:

self.count = 0

self.max_count stores the best frequency found so far:

self.max_count = 0

The update step compares the current node with the previous node:

if self.prev and self.prev.val == node.val:
    self.count += 1
else:
    self.count = 1

If the current streak is better than the previous best, the old answer is no longer valid:

if self.count > self.max_count:
    self.max_count = self.count
    self.ans = [node.val]

If the current streak ties the best frequency, the value is also a mode:

elif self.count == self.max_count:
    self.ans.append(node.val)

Finally, the current node becomes the previous node for the next inorder visit:

self.prev = node

Testing

def sorted_modes(root):
    return sorted(Solution().findMode(root))

def test_find_mode():
    # Tree:
    #   1
    #    \
    #     2
    #    /
    #   2
    root = TreeNode(1)
    root.right = TreeNode(2)
    root.right.left = TreeNode(2)
    assert sorted_modes(root) == [2]

    # Tree:
    #     2
    #    / \
    #   1   3
    root = TreeNode(2, TreeNode(1), TreeNode(3))
    assert sorted_modes(root) == [1, 2, 3]

    # Tree:
    #     2
    #    / \
    #   2   3
    #        \
    #         3
    root = TreeNode(2)
    root.left = TreeNode(2)
    root.right = TreeNode(3)
    root.right.right = TreeNode(3)
    assert sorted_modes(root) == [2, 3]

    # Single node
    root = TreeNode(10)
    assert sorted_modes(root) == [10]

    print("all tests passed")

Test meaning: