leetcode@ [116/117] Populating Next Right Pointers in Each Node I & II (Tree, BFS)

Posted 2020-06-21 流白

https://leetcode.com/problems/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

/**
 * 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) {}
 * };
 */
class node {
    public:
        TreeLinkNode* nd; 
        int lv;
        node(TreeLinkNode* rhs, int l): nd(rhs), lv(l) {}
};

class Solution {
public:
    void connect(TreeLinkNode *root) {
        if(!root)  return;
        
        queue<node> q;
        q.push(node(root, 0));
        vector<node> vec;
        
        while(!q.empty()) {
            node top = q.front();
            vec.push_back(top);
            q.pop();
            
            int cur_lv = top.lv;
            TreeLinkNode* cur_nd = top.nd;
            
            if(cur_nd->left)  q.push(node(cur_nd->left, cur_lv+1));
            if(cur_nd->right)  q.push(node(cur_nd->right, cur_lv+1));
        }
        
        int l = 0, r = 1;
        while(r < vec.size()) {
            vec[l].nd->next = NULL;
            while(r < vec.size() && vec[r].lv == vec[l].lv) {
                vec[l].nd->next = vec[r].nd;
                ++l;
                ++r;
            }
            
            l = r;
            r = l+1;
        }
        vec[vec.size()-1].nd->next = NULL;
    }
};