diff --git a/problems/0102.二叉树的层序遍历.md b/problems/0102.二叉树的层序遍历.md index 6235ad21..47969b25 100644 --- a/problems/0102.二叉树的层序遍历.md +++ b/problems/0102.二叉树的层序遍历.md @@ -2096,36 +2096,40 @@ class Solution { python代码: ```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: 'Node') -> 'Node': if not root: - return None - queue = [root] + return root + + queue = collections.deque([root]) + while queue: - n = len(queue) - for i in range(n): - node = queue.pop(0) + level_size = len(queue) + prev = None + + for i in range(level_size): + node = queue.popleft() + + if prev: + prev.next = node + + prev = node + if node.left: queue.append(node.left) + if node.right: queue.append(node.right) - if i == n - 1: - break - node.next = queue[0] - return root - -# 链表解法 -class Solution: - def connect(self, root: 'Node') -> 'Node': - first = root - while first: - cur = first - while cur: # 遍历每一层的节点 - if cur.left: cur.left.next = cur.right # 找左节点的next - if cur.right and cur.next: cur.right.next = cur.next.left # 找右节点的next - cur = cur.next # cur同层移动到下一节点 - first = first.left # 从本层扩展到下一层 + return root ```