This problem is solved with the idea of BFS. Due to the exist of the "next" pointer, after each level of the binary tree is traveled, this level generates a linked list, thus, the variable needs to be stored is just the head node of this linked list.
class Solution: # @param root, a tree link node # @return nothing def connect(self, root): if root is None: return nd1 = root nd2 = nd1.left while nd2 is not None: nd2.next = nd1.right nd3 = nd2.next nd1 = nd1.next while nd1 is not None: nd3.next = nd1.left nd3 = nd3.next nd3.next = nd1.right nd3 = nd3.next nd1 = nd1.next nd1 = nd2 nd2 = nd2.left
It can be shorted to something like the following:
class Solution: # @param root, a tree link node # @return nothing def connect(self, root): head = root while head: # connect one level start = head while start: try: # make left child's next point to right child start.left.next = start.right # do cross subtree connection start.right.next = start.next.left start = start.next except AttributeError: break head = head.left