116. Populating Next Right Pointers in Each Node

Information

  • Diffculty: Medium

  • Created: 2020-01-12 11:52:20

  • Last Motified: 2020-01-12 11:52:20

  • Tags: Tree, Depth-first Search

Intuition

This problem could be solved by BFS. We mark the level for each node, and for the node with the same level as previous one, the next field of previous node must be the current node.

Solution

from collections import deque


# 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: 'Node') -> 'Node':
        last = None
        last_level = -1
        queue = deque([[root, 0]])

        while queue:
            node, level = queue.popleft()
            if node is None:
                continue
            if last_level == level:
                last.next = node
            last, last_level = node, level
            queue.append([node.left, level + 1])
            queue.append([node.right, level + 1])
        return root