LeetCode 117 - 填充每个节点的下一个右侧节点指针 II

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

def create_perfect_binary_tree(nums):
    """
    创建一个包含给定数字的完美二叉树。

    Args:
        nums: 一个包含二叉树节点值的列表。

    Returns:
        完美二叉树的根节点。
    """

    if not nums:
        return None

    # 计算树的深度
    depth = 0
    num_nodes = len(nums)
    while (1 << (depth + 1)) - 1 < num_nodes:
        depth += 1

    def build_tree(index, level):
        """
        递归构建二叉树。

        Args:
            index: 当前节点在列表 nums 中的索引。
            level: 当前节点的深度。

        Returns:
            当前节点。
        """

        if level > depth or index >= num_nodes:
            return None

        node = Node(nums[index])

        node.left = build_tree(2 * index + 1, level + 1)
        node.right = build_tree(2 * index + 2, level + 1)

        return node

    return build_tree(0, 0)

# 测试代码
nums = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]
root = create_perfect_binary_tree(nums)

# 验证二叉树是否创建成功（可选）
def print_tree(node):
    if node:
        print(node.val)
        if node.next:
                print(node.val, node.next.val, sep="->")
        print_tree(node.left)
        print_tree(node.right)

class Solution:
    q = {}
    def connectIMPL(self, root: 'Optional[Node]', depth: int) -> 'Optional[Node]':
        if root is None:
            return None
        # 左子节点只要存在，就必须执行next指向右子节点，即便右子节点为空
        if root.left:
            root.left.next = root.right
            if self.q:
                item = self.q.get(depth+1)
                if item:
                    lastRight = self.q.pop(depth + 1)
                    lastRight.next = root.left
            if root.right:
                self.q[depth+1]=root.right
            else:
                # 此处和上一题不同，如果右子节点为空，则左子节点进入字典，因为左子节点将代替不存在的右子节点继续连接同级的右边的结点
                self.q[depth+1]= root.left
        elif not root.left and root.right:
        # 如果左子节点为空，则需要从字典读出同级结点进行连接
            if self.q:
                item = self.q.get(depth+1)
                if item:
                    lastRight = self.q.pop(depth + 1)
                    lastRight.next = root.right
            self.q[depth+1] = root.right
            
        self.connectIMPL(root.left, depth + 1)
        self.connectIMPL(root.right, depth + 1)
        return root
    def connect(self, root: 'Optional[Node]') -> 'Optional[Node]':
        return self.connectIMPL(root, 0)

s = Solution()
s.connect(root)
print_tree(root)