1def connect(root):
2 if not root:
3 return root
4
5 queue = [root]
6 level = 0
7
8 while queue:
9 level_size = len(queue)
10
11 for i in range(level_size):
12 current = queue.pop(0)
13
14 # Connect to next node in level
15 if i < level_size - 1:
16 current.next = queue[0]
17
18 # Add children to queue
19 if current.left:
20 queue.append(current.left)
21 if current.right:
22 queue.append(current.right)
23
24 level += 1
25
26 return root