[LeetCode] 117. Populating Next Right Pointers in Each Node II

Problem : https://leetcode.com/problems/populating-next-right-pointers-in-each-node-ii/

Be careful to link to the leftmost node of right sub-trees. 

"""
# 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':
        
        def right_node():
            p = root.next
            
            # find the leftmost node of the right subtrees
            while p:
                if p.left:
                    return p.left
                if p.right:
                    return p.right
                
                p = p.next
                
            return None
        
        if root:
            if root.left:
                if root.right:
                    root.left.next = root.right
                else:
                    root.left.next = right_node()
                
            if root.right:
                root.right.next = right_node()
            
            # should link next level right sub-trees first
            # to guarantee nodes on left sub-tree can always
            # find the leftmost node on right sub-trees
            self.connect(root.right)  
            self.connect(root.left)
        
        return root