Update 0102.二叉树的层序遍历.md

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]
-        while queue:
-            n = len(queue)
-            for i in range(n):
-                node = queue.pop(0)
-                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  # 从本层扩展到下一层
+        queue = collections.deque([root])
+        
+        while queue:
+            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)
+        
         return root
 ```