LeetCode 998: Maximum Binary Tree II

Problem Restatement

We are given the root of a maximum binary tree and an integer val.

The tree was originally built from an array a using this rule:

1. The largest value in a becomes the root.
2. The left subtree is built from the elements to the left of that largest value.
3. The right subtree is built from the elements to the right of that largest value.

Now we append val to the end of the original array and rebuild the maximum binary tree.

We need to return the new root.

The important detail is that val is appended to the end of the array, not inserted anywhere else. Source: LeetCode 998 problem statement.

Input and Output

Function shape:

def insertIntoMaxTree(root: Optional[TreeNode], val: int) -> Optional[TreeNode]:
    ...

Examples

Example 1:

root = [4, 1, 3, None, None, 2]
val = 5

The new value 5 is greater than the current root 4.

So 5 becomes the new root, and the old tree becomes its left subtree.

Answer:

[5, 4, None, 1, 3, None, None, 2]

Example 2:

root = [5, 2, 4, None, 1]
val = 3

The new value 3 is less than the root 5, so it belongs somewhere in the right subtree.

It is greater than the node 4’s right-side position where it should attach, so it becomes a node on the right spine.

First Thought: Rebuild the Whole Tree

One direct solution is:

1. Recover the original array from the tree.
2. Append val.
3. Rebuild the maximum binary tree from the new array.

This works, but it does unnecessary work.

The problem gives us the tree, and val is appended to the end of the original array. That means only the right side of the tree can change.

Key Insight

Because val is appended to the end of the original array, it is to the right of every existing value.

In a maximum binary tree, the right subtree represents values to the right of the root in the original array.

So insertion only affects the right spine of the tree.

There are two cases:

1. If val > root.val, then val is the largest value in the new array. It becomes the new root, and the old tree becomes its left child.

2. Otherwise, the root stays the same. Since val is appended to the end, it must be inserted into the right subtree.

This gives a simple recursive rule.

Algorithm

For a node root:

1. If root is None, return a new node with val.
2. If val > root.val:
   • create a new node with value val
   • set its left child to root
   • return the new node
3. Otherwise:
   • recursively insert val into root.right
   • return root

Correctness

If root is None, the array segment is empty, so appending val creates a single-node tree.

If val > root.val, then val is larger than every value in the current maximum tree, because root.val is the maximum value of that tree. Since val is appended after all old values, all old values lie to its left in the new array. Therefore, val must become the root, with the old tree as its left subtree.

If val < root.val, then the original root remains the largest value in this array segment, so the root does not change. Since val was appended to the end, it belongs in the portion of the array to the right of root, which corresponds to root.right. Recursively inserting into root.right therefore constructs the correct new right subtree.

By applying this reasoning recursively, the algorithm returns exactly the maximum binary tree after appending val.

Complexity

Let h be the height of the tree.

In the worst case, h can be n.

Implementation

# 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 insertIntoMaxTree(
        self,
        root: Optional[TreeNode],
        val: int,
    ) -> Optional[TreeNode]:
        if root is None:
            return TreeNode(val)

        if val > root.val:
            return TreeNode(val, root, None)

        root.right = self.insertIntoMaxTree(root.right, val)
        return root

Code Explanation

If the current subtree is empty, the inserted value becomes a new node:

if root is None:
    return TreeNode(val)

If val is larger than the current root, it becomes the root of this subtree:

if val > root.val:
    return TreeNode(val, root, None)

The old subtree becomes the left child because all old values are before val in the array.

Otherwise, the current root stays in place, and insertion continues into the right subtree:

root.right = self.insertIntoMaxTree(root.right, val)
return root

Testing

from typing import Optional

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def preorder(root):
    if root is None:
        return []
    return [root.val] + preorder(root.left) + preorder(root.right)

def run_tests():
    s = Solution()

    root1 = TreeNode(
        4,
        TreeNode(1),
        TreeNode(3, TreeNode(2)),
    )
    new_root1 = s.insertIntoMaxTree(root1, 5)
    assert preorder(new_root1) == [5, 4, 1, 3, 2]

    root2 = TreeNode(
        5,
        TreeNode(2, None, TreeNode(1)),
        TreeNode(4),
    )
    new_root2 = s.insertIntoMaxTree(root2, 3)
    assert preorder(new_root2) == [5, 2, 1, 4, 3]

    root3 = TreeNode(5)
    new_root3 = s.insertIntoMaxTree(root3, 1)
    assert preorder(new_root3) == [5, 1]

    root4 = TreeNode(1)
    new_root4 = s.insertIntoMaxTree(root4, 2)
    assert preorder(new_root4) == [2, 1]

    print("all tests passed")

run_tests()