116. 填充每个节点的下一个右侧节点指针

原题

https://leetcode.cn/problems/populating-next-right-pointers-in-each-node/description/

思路

广度优先搜索

复杂度

时间：O(n)
空间：O(n)

Python代码

"""
# Definition for a Node.
class Node:
    def __init__(self, val: int = 0, left: 'Node' = None, right: 'Node' = None, next: 'Node' = None):
        self.val = val
        self.left = left
        self.right = right
        self.next = next
"""

class Solution:
    def connect(self, root: 'Optional[Node]') -> 'Optional[Node]':
        if not root:
            return None 
        q = [root]
        while q:
            if len(q) > 1:
                for i in range(len(q)- 1):
                    q[i].next = q[i+1]
            for i in range(len(q)):
                node = q.pop(0)
                if node.left:
                    q.append(node.left)
                if node.right:
                    q.append(node.right)

        return root
        

Go代码

/**
 * Definition for a Node.
 * type Node struct {
 *     Val int
 *     Left *Node
 *     Right *Node
 *     Next *Node
 * }
 */

func connect(root *Node) *Node {
	if root == nil {
		return nil
	}
	q := []*Node{root}
	for len(q) > 0 {
		n := len(q)
		// 指针指向其下一个右侧节点
		if len(q) > 1 {
			for i := 0; i < n-1; i++ {
				q[i].Next = q[i+1]
			}
		}
		// 将下一层节点加入队列
		for i := 0; i < n; i++ {
			node := q[0]
			q = q[1:]
			if node.Left != nil {
				q = append(q, node.Left)
			}
			if node.Right != nil {
				q = append(q, node.Right)
			}
		}
	}
	return root
}