N叉树遍历模板

taixian / 2024-02-18 / 原文

N叉树(N-ary Tree)的类型和代码模板与二叉树有些相似,但由于N叉树具有多个子节点,因此在遍历和节点定义上有所不同。以下是N叉树的类型和相应的代码模板:

N叉树节点的定义:

class Node:
    def __init__(self, val=None, children=None):
        self.val = val
        self.children = children if children is not None else []

N叉树的遍历(递归和迭代):

1. 前序遍历(Preorder Traversal)

递归实现:

def preorder(root):
    result = []
    if root:
        result.append(root.val)
        for child in root.children:
            result.extend(preorder(child))
    return result

例子:N叉树前序遍历

class Solution:
    def preorder(self, root: 'Node') -> List[int]:
        res = []
        def dfs(root):
            if root is None:
                return
            res.append(root.val)
            for i in root.children:
                dfs(i)
        dfs(root)
        return res

迭代实现:

def preorder(root):
    result = []
    q = [root]
    while q:
        node = q.pop()
        if node:
            result.append(node.val)
            q.extend(reversed(node.children)) #或者使用 for child in node.children[::-1]
    return result

注意: reversed() 是 Python 内置函数,用于反转一个可迭代对象的元素顺序。在这个上下文中,node.children 是一个子节点列表,reversed(node.children) 返回一个反转后的子节点列表。在遍历 N 叉树时,我们使用栈来模拟递归调用。由于栈是先进后出的数据结构,我们希望先处理的子节点后进栈,以保证先处理左边的子节点。因此,我们使用 reversed() 来反转子节点列表,使得先处理的子节点在栈顶。
具体而言,stack.extend(reversed(node.children)) 将 node.children 中的子节点按照逆序添加到栈中。这样,下一次从栈中弹出的节点就是先处理的子节点,符合前序遍历的顺序。

2. 后序遍历(Postorder Traversal)

递归实现:

def postorder(root):
    result = []
    if root:
        for child in root.children:
            result.extend(postorder(child))
        result.append(root.val)
    return result

迭代实现:

def postorder(root):
    result = []
    stack = [root]
    while stack:
        node = stack.pop()
        if node:
            result.append(node.val)
            stack.extend(node.children)
    return result[::-1]

对于N叉树的层序遍历(Level Order Traversal),我们可以使用队列来实现。以下是N叉树层序遍历的代码模板

"""
# Definition for a Node.
class Node:
    def __init__(self, val=None, children=None):
        self.val = val
        self.children = children
"""

class Solution:
    def levelOrder(self, root: 'Node') -> List[List[int]]:
        if root is None:
            return 
        q = deque()
        q.append(root)
        res = []
        while q:
            vals = []
            for _ in range(len(q)):
                node = q.popleft()
                vals.append(node.val) #node.children 返回的是一个子节点列表
                #extend 可以将 node.children 中的每个元素添加到队列,
                #而不是将整个列表作为一个元素添加
                if node.children: q.extend(node.children)
            res.append(vals)
        return res