Given a binary tree with each node having one extra field of nextRight. nextRight points to the next right node of the same level.
Initially all nodes‘ nextRight field are set to null. Connect all nodes at the same level by setting each node‘s nextRight field.
Solution 1. O(n) runtime, O(n) space, level order traversal
This algorithm is the most straightforward solution.
1 class SpecialTreeNode {
2 int val;
3 SpecialTreeNode left;
4 SpecialTreeNode right;
5 SpecialTreeNode nextRight;
6 SpecialTreeNode(int v) {
7 val = v;
8 left = null;
9 right = null;
10 nextRight = null;
11 }
12 }
13 public class ConnectNodesOfSameLevel {
14
15 public static void connectNodesLevelOrder(SpecialTreeNode root) {
16 if(root == null) {
17 return;
18 }
19 Queue<SpecialTreeNode> queue = new LinkedList<SpecialTreeNode>();
20 queue.add(root);
21 SpecialTreeNode prev = null, curr = null;
22
23 while(!queue.isEmpty()) {
24 prev = null;
25 int size = queue.size();
26 for(int i = 0; i < size; i++) {
27 curr = queue.poll();
28 if(prev != null) {
29 prev.nextRight = curr;
30 }
31 if(curr.left != null) {
32 queue.add(curr.left);
33 }
34 if(curr.right != null) {
35 queue.add(curr.right);
36 }
37 prev = curr;
38 }
39 }
40 }
41 }
Solution 2. O(n) runtime, O(height + width) space for recursion stack, a variation of pre-order traversal.
1. for the current node A, recursively set the left/right child nodes‘ nextRight of other nodes at the same level with A.
This is done by recursively calling the connect method on A.nextRight.
2. then set A‘s left/right child nodes‘ nextRight. Since A‘s nextRight has already been set when processing the previous level,
we use A‘s nextRight to find the next right node for A‘s left/right child nodes at A‘s children nodes‘ level.
1 public static void connectNodesRecursion(SpecialTreeNode root) {
2 if(root == null) {
3 return;
4 }
5 //set the nextRight of children nodes of other nodes at the same level
6 connectNodesRecursion(root.nextRight);
7
8 //set the nextRight of the current node‘s left/right child node
9 if(root.left != null && root.right != null) {
10 root.left.nextRight = root.right;
11 root.right.nextRight = getNextRightNodeOfSameLevel(root);
12 connectNodesRecursion(root.left);
13 }
14 else if(root.left != null) {
15 root.left.nextRight = getNextRightNodeOfSameLevel(root);
16 connectNodesRecursion(root.left);
17 }
18 else if(root.right != null) {
19 root.right.nextRight = getNextRightNodeOfSameLevel(root);
20 connectNodesRecursion(root.right);
21 }
22 }
23 private static SpecialTreeNode getNextRightNodeOfSameLevel(SpecialTreeNode node) {
24 if(node == null) {
25 return null;
26 }
27 SpecialTreeNode temp = node.nextRight;
28 while(temp != null && temp.left == null && temp.right == null) {
29 temp = temp.nextRight;
30 }
31 if(temp == null) {
32 return null;
33 }
34 else if(temp.left != null) {
35 return temp.left;
36 }
37 return temp.right;
38 }
Solution 3. O(n) runtime, O(1) space, iterative implementation of solution 2.
The outer loop goes through all nodes‘ on the left edge of the given binary tree.
For each node A that the outer loop iterates, the inner loop goes through all nodes of the same level with A and
sets each node‘s children nodes‘ nextRight field.
The core idea of both solution 2 and 3 is that when processing nodes at level k, their children nodes at level k + 1
are updated with correct nextRight value. The algorithm is designed this way is because without using level order
traversal, the only way of setting nextRight values level by level is to use the already set nextRight values from the
previous level nodes.
1 public static void connectNodesIteration(SpecialTreeNode root) {
2 if(root == null) {
3 return;
4 }
5 SpecialTreeNode leftMost = root, curr = null;
6 leftMost.nextRight = null;
7
8 while(leftMost != null) {
9 curr = leftMost;
10 while(curr != null) {
11 if(curr.left != null && curr.right != null) {
12 curr.left.nextRight = curr.right;
13 curr.right.nextRight = getNextRightNodeOfSameLevel(curr);
14 }
15 else if(curr.left != null) {
16 curr.left.nextRight = getNextRightNodeOfSameLevel(curr);
17 }
18 else if(curr.right != null) {
19 curr.right.nextRight = getNextRightNodeOfSameLevel(curr);
20 }
21 curr = curr.nextRight;
22 }
23 leftMost = leftMost.left;
24 }
25 }