I and two can use exactly same code.
1 /** 2 * Definition for binary tree with next pointer. 3 * struct TreeLinkNode { 4 * int val; 5 * TreeLinkNode *left, *right, *next; 6 * TreeLinkNode(int x) : val(x), left(NULL), right(NULL), next(NULL) {} 7 * }; 8 */ 9 class Solution { 10 public: 11 TreeLinkNode *getNext(TreeLinkNode *root) { 12 if (!root) return NULL; 13 if (root->left) return root->left; 14 if (root->right) return root->right; 15 return getNext(root->next); 16 } 17 void connect(TreeLinkNode *root) { 18 if (!root) return; 19 TreeLinkNode *current = root; 20 while (current) { 21 root = current; 22 while (root) { 23 if (root->left && root->right) { 24 root->left->next = root->right; 25 root->right->next = getNext(root->next); 26 } else if (root->left || root->right) { 27 getNext(root)->next = getNext(root->next); 28 } 29 root = root->next; 30 } 31 current = getNext(current); 32 } 33 } 34 };