Leetcode 117. Populating Next Right Pointers in Each Node II

Problem

Given a binary tree

struct Node {
  int val;
  Node *left;
  Node *right;
  Node *next;
}

Populate each next pointer to point to its next right node. If there is no next right node, the next pointer should be set to NULL.

Initially, all next pointers are set to NULL.

Algorithm

Recursion. Use dfs to visit the tree and save the node in different lists, then set the next point with the lists.

Code

"""
# 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':
        
        Q = []
        def dfs(root, d):
            if not root:
                return
            if len(Q) <= d:
                Q.append([root])
            else:
                Q[d].append(root)
                
            dfs(root.left, d+1)
            dfs(root.right, d+1)
            
        dfs(root, 0) 
        
        for i in range(len(Q)):
            for j in range(len(Q[i])-1):
                Q[i][j].next = Q[i][j+1]
                
        return root