【Leetcode】117. Populating Next Right Pointers in Each Node II

题目地址：

https://leetcode.com/problems/populating-next-right-pointers-in-each-node-ii/

给定一棵二叉树， 每个节点除了有两个指针指向左右孩子以外，还有一个next指针。要求修改该二叉树，使得每个节点的next指针指向其同一层的下一个右边的节点，如果右边没有则指向null。返回修改后的树根。

思路是BFS。开一个队列对树进行BFS。分层遍历，当遇到每一层第一个节点时，用一个变量保存起来。接着遍历同层下一个节点的时候，用保存起来的那个节点的next指向这个节点，并更新它为当前节点即可。代码如下：

/*
// Definition for a Node.
class Node {
public:
    int val;
    Node* left;
    Node* right;
    Node* next;

    Node() : val(0), left(NULL), right(NULL), next(NULL) {}

    Node(int _val) : val(_val), left(NULL), right(NULL), next(NULL) {}

    Node(int _val, Node* _left, Node* _right, Node* _next)
        : val(_val), left(_left), right(_right), next(_next) {}
};
*/
class Solution {
 public:
  Node* connect(Node* root) {
    if (!root) return nullptr;
    queue<Node*> q;
    q.push(root);

    Node* prev = nullptr;
    while (q.size()) {
      for (int i = q.size(); i; i--) {
        auto cur = q.front();
        q.pop();
        if (cur->left) q.push(cur->left);
        if (cur->right) q.push(cur->right);
        // 如果是第一个节点，保存起来；
        // 否则将next指向当前节点并更新prev
        if (!prev)
          prev = cur;
        else {
          prev->next = cur;
          prev = cur;
        }
      }

      prev = nullptr;
    }
    return root;
  }
};

时空复杂度$O(n)$。