Populating Next Right Pointers in Each Node II 解答

 1 /**
 2  * Definition for binary tree with next pointer.
 3  * public class TreeLinkNode {
 4  *     int val;
 5  *     TreeLinkNode left, right, next;
 6  *     TreeLinkNode(int x) { val = x; }
 7  * }
 8  */
 9 public class Solution {
10     public void connect(TreeLinkNode root) {
11         TreeLinkNode prevHead = root, prevCur = root;
12         TreeLinkNode currentHead = null, current = null;
13         while (prevHead != null) {
14             prevCur = prevHead;
15             // Find current head
16             while (prevCur != null && prevCur.left == null && prevCur.right == null) {
17                 prevCur = prevCur.next;
18             }
19             if (prevCur == null) {
20                 break;
21             } else if (prevCur.left != null) {
22                 currentHead = prevCur.left;
23                 current = currentHead;
24                 if (prevCur.right != null) {
25                     current.next = prevCur.right;
26                     current = current.next;
27                 }
28             } else {
29                 currentHead = prevCur.right;
30                 current = currentHead;
31             }
32             // link current level
33             prevCur = prevCur.next;
34             while (prevCur != null) {
35                 if (prevCur.left != null) {
36                     current.next = prevCur.left;
37                     current = current.next;
38                 }
39                 if (prevCur.right != null) {
40                     current.next = prevCur.right;
41                     current = current.next;
42                 }
43                 prevCur = prevCur.next;
44             }
45             // reset
46             prevHead = currentHead;
47         }
48         
49     }
50 }