LeetCode 117. Populating Next Right Pointers in Each Node II

题目：

Follow up for problem "Populating Next Right Pointers in Each Node".

What if the given tree could be any binary tree? Would your previous solution still work?

Note:

• You may only use constant extra space.

 

For example,
Given the following binary tree,

         1
       /  \
      2    3
     / \    \
    4   5    7

 

After calling your function, the tree should look like:

         1 -> NULL
       /  \
      2 -> 3 -> NULL
     / \    \
    4-> 5 -> 7 -> NULL

思路： 类似宽度优先搜索， 保存下一层的起始节点。  注意当每次处理一个结点的时候，如果父亲的next没有需要继续往右找。
Wa了好几次。 继续加油。 写代码需要再认真一点，写错好几个地方。

/**
 * Definition for binary tree with next pointer.
 * struct TreeLinkNode {
 *  int val;
 *  TreeLinkNode *left, *right, *next;
 *  TreeLinkNode(int x) : val(x), left(NULL), right(NULL), next(NULL) {}
 * };
 */

#include <vector>
using namespace std;

class Solution {
public:
    void connect(TreeLinkNode *root) {
        if(root == NULL) return ;
        TreeLinkNode* ptr = root;
        TreeLinkNode* wait = NULL;
        while(ptr || wait)
        {
            if(ptr == NULL)
            {
                ptr = wait;
                wait = NULL;
            }
            if(wait == NULL && ptr->left)  wait = ptr->left;
            if(wait == NULL && ptr->right) wait = ptr->right;
            
            vector<TreeLinkNode*> lk;
            if(ptr->left)  lk.push_back(ptr->left);
            if(ptr->right) lk.push_back(ptr->right);
            ptr = ptr->next;
            if(lk.size() == 0)
            {
                continue;
            }
            
            while(ptr && ptr->left==NULL && ptr->right==NULL) ptr = ptr->next;
            if(ptr)
            {
                if(ptr->left) lk.push_back(ptr->left);
                else if(ptr->right) lk.push_back(ptr->right);
            }
            int num = lk.size();
            for(int i=0; i<num-1; i++)
            {
                lk[i]->next = lk[i+1];
            }
        }
    }
};