diff --git a/src/populating_next_right_pointers_in_each_node_ii.py b/src/populating_next_right_pointers_in_each_node_ii.py new file mode 100644 index 0000000..d22c3ef --- /dev/null +++ b/src/populating_next_right_pointers_in_each_node_ii.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from collections import deque + + +class Node: + def __init__( + self, + val: int = 0, + left: Node | None = None, + right: Node | None = None, + next: Node | None = None, + ): + self.val = val + self.left = left + self.right = right + self.next = next + + +class Solution: + def connect(self, root: Node | None) -> Node | None: + if not root: + return None + + queue: deque[Node] = deque() + queue.append(root) + + while queue: + for i in range(len(queue) - 1): + queue[i].next = queue[i + 1] + for _ in range(len(queue)): + node = queue.popleft() + if node.left: + queue.append(node.left) + if node.right: + queue.append(node.right) + + return root diff --git a/tests/test_populating_next_right_pointers_in_each_node_ii.py b/tests/test_populating_next_right_pointers_in_each_node_ii.py new file mode 100644 index 0000000..d07cf29 --- /dev/null +++ b/tests/test_populating_next_right_pointers_in_each_node_ii.py @@ -0,0 +1,21 @@ +from src.populating_next_right_pointers_in_each_node_ii import Node +from src.populating_next_right_pointers_in_each_node_ii import Solution + + +def test_example(): + n7 = Node(7) + n3 = Node(3, right=n7) + n4 = Node(4) + n5 = Node(5) + n2 = Node(2, left=n4, right=n5) + n1 = Node(1, left=n2, right=n3) + + Solution().connect(n1) + + assert n1.next is None + assert n2.next is n3 and n3.next is None + assert n4.next is n5 and n5.next is n7 and n7.next is None + + +def test_empty(): + assert Solution().connect(None) is None