116. Populating Next Right Pointers in Each Node

题目

Given a binary tree

struct TreeLinkNode {
  TreeLinkNode *left;
  TreeLinkNode *right;
  TreeLinkNode *next;
}

Populate each next pointer to point to its next right node. If there is no next right node, the next pointer should be set to NULL.

Initially, all next pointers are set to NULL.

Note:

You may only use constant extra space.
You may assume that it is a perfect binary tree (ie, all leaves are at the same level, and every parent has two children).
For example,
Given the following perfect binary tree,

After calling your function, the tree should look like:

分析

题目要求，将叶子结点的next指向它的右边的结点。如果右边结点不存在，则为NULL.

思路

从题目看来，我们可以层序遍历，然后标记结尾，对结尾做特殊处理。
这里我们对层序遍历做一些修改。因为二叉树为完全二叉树，所以每一次的个数为2^n个。所以只需要记录层数即可。每层访问了2^n即到达了最后一个结点。

代码

class Solution {
    public:
        void connect(TreeLinkNode *root) {
            if(root == NULL){
                return;
            }
            queue<TreeLinkNode*> t_queue;
            t_queue.push(root);
            int flag = 1;
            int time = 0;
            TreeLinkNode* pre = NULL;
            while(!t_queue.empty()){
                TreeLinkNode *node = t_queue.front();
                time ++;

                if(time == 1){
                    pre = node;

                }else{
                    pre->next = node;
                    pre = node;
                }

                if(time == flag){
                    node->next = NULL;
                    flag = flag << 1;
                    time = 0;
                }

                if(node->left != NULL){
                    t_queue.push(node->left);
                }
                  if(node->right != NULL){
                    t_queue.push(node->right);
                }

                t_queue.pop();
            }
        }
    };