LeetCode 116: Populating Next Right Pointers in Each Node

Problem Restatement

We are given a perfect binary tree.

A perfect binary tree has two important properties:

1. Every internal node has exactly two children.
2. All leaves are on the same level.

Each node has four fields:

val
left
right
next

We need to populate each next pointer so that it points to the node immediately to its right on the same level. If there is no node to the right, next should be None. Initially, all next pointers are None. The official statement also asks for constant extra space, with recursive stack space allowed.

For this tree:

        1
      /   \
     2     3
    / \   / \
   4   5 6   7

After connecting:

        1 -> None
      /   \
     2  -> 3 -> None
    / \   / \
   4 ->5->6->7 -> None

Input and Output

LeetCode gives this node class:

class Node:
    def __init__(self, val: int = 0, left: 'Node' = None, right: 'Node' = None, next: 'Node' = None):
        self.val = val
        self.left = left
        self.right = right
        self.next = next

The function shape is:

class Solution:
    def connect(self, root: 'Optional[Node]') -> 'Optional[Node]':
        ...

Examples

Consider:

        1
      /   \
     2     3
    / \   / \
   4   5 6   7

At level 0, only node 1 exists:

1 -> None

At level 1, node 2 is followed by node 3:

2 -> 3 -> None

At level 2, the nodes are connected from left to right:

4 -> 5 -> 6 -> 7 -> None

The returned tree root is still node 1, but its next pointers have been filled.

For an empty tree:

root = None

the answer is:

None

First Thought: Use Level Order Traversal

A simple solution is breadth-first search.

For each level, we can store nodes in a queue and connect each node to the next node in that level.

That works, but it uses extra queue space.

The follow-up asks us to use constant extra space. Since the tree is perfect, we can do better.

Key Insight

Because the tree is perfect, every node with children has both a left child and a right child.

For any node cur, there are two kinds of connections we need to make.

First, connect its own children:

cur.left.next = cur.right

Second, connect across neighboring parents:

cur.right.next = cur.next.left

The second connection works only when cur.next exists.

For example:

        1
      /   \
     2  -> 3
    / \   / \
   4   5 6   7

At node 2:

2.left.next = 2.right

connects:

4 -> 5

And:

2.right.next = 2.next.left

connects:

5 -> 6

That is the cross-parent connection.

Algorithm

Use the already-built next pointers to walk level by level.

Start with:

leftmost = root

leftmost points to the first node of the current level.

While leftmost.left exists:

1. Set cur = leftmost.
2. Walk across the current level using cur.next.
3. For each cur:
   1. Connect cur.left.next = cur.right.
   2. If cur.next exists, connect cur.right.next = cur.next.left.
   3. Move cur = cur.next.
4. Move down to the next level with leftmost = leftmost.left.

Return root.

Correctness

The algorithm processes the tree level by level, starting at the root.

At any node cur, the connection cur.left.next = cur.right correctly links the left child to the right child under the same parent.

If cur.next exists, then cur has a neighboring parent to its right on the same level. Since the tree is perfect, that neighboring parent has a left child. The next node to the right of cur.right is exactly cur.next.left, so cur.right.next = cur.next.left is correct.

These two assignments cover every horizontal connection on the next level:

• Connections inside the same parent.
• Connections between two neighboring parents.

The current level can be traversed using next pointers that were already created while processing the previous level. The root level needs no connection, so the process can start there.

By induction over the levels, after processing a level, all next pointers on the level below it are correct. Therefore, after the loop finishes, all levels have correct next pointers.

Complexity

Here, n is the number of nodes in the tree.

The output pointers are part of the given nodes, so they do not count as extra space.

Implementation

from typing import Optional

# Definition for a Node.
# class Node:
#     def __init__(self, val: int = 0, left: 'Node' = None, right: 'Node' = None, next: 'Node' = None):
#         self.val = val
#         self.left = left
#         self.right = right
#         self.next = next

class Solution:
    def connect(self, root: 'Optional[Node]') -> 'Optional[Node]':
        if root is None:
            return None

        leftmost = root

        while leftmost.left is not None:
            cur = leftmost

            while cur is not None:
                cur.left.next = cur.right

                if cur.next is not None:
                    cur.right.next = cur.next.left

                cur = cur.next

            leftmost = leftmost.left

        return root

Code Explanation

Handle the empty tree:

if root is None:
    return None

Start at the leftmost node of the current level:

leftmost = root

The loop continues while there is a lower level to connect:

while leftmost.left is not None:

For each level, walk from left to right:

cur = leftmost

while cur is not None:

Connect the two children of the same parent:

cur.left.next = cur.right

Then connect across parents if a neighboring parent exists:

if cur.next is not None:
    cur.right.next = cur.next.left

Move to the next parent on the same level:

cur = cur.next

After finishing the level, move down:

leftmost = leftmost.left

Return the original root:

return root

Testing

from typing import Optional

class Node:
    def __init__(self, val: int = 0, left: 'Node' = None, right: 'Node' = None, next: 'Node' = None):
        self.val = val
        self.left = left
        self.right = right
        self.next = next

class Solution:
    def connect(self, root: Optional[Node]) -> Optional[Node]:
        if root is None:
            return None

        leftmost = root

        while leftmost.left is not None:
            cur = leftmost

            while cur is not None:
                cur.left.next = cur.right

                if cur.next is not None:
                    cur.right.next = cur.next.left

                cur = cur.next

            leftmost = leftmost.left

        return root

def levels_by_next(root):
    ans = []
    leftmost = root

    while leftmost is not None:
        level = []
        cur = leftmost

        while cur is not None:
            level.append(cur.val)
            cur = cur.next

        ans.append(level)
        leftmost = leftmost.left

    return ans

def run_tests():
    s = Solution()

    root1 = Node(
        1,
        Node(2, Node(4), Node(5)),
        Node(3, Node(6), Node(7)),
    )

    s.connect(root1)

    assert levels_by_next(root1) == [
        [1],
        [2, 3],
        [4, 5, 6, 7],
    ]

    root2 = Node(1)

    s.connect(root2)

    assert levels_by_next(root2) == [[1]]

    root3 = None

    assert s.connect(root3) is None

    print("all tests passed")

run_tests()

Test meaning: