如何实现基于指定深度拆分树分支的split_tree函数？

问题分析 #

核心需求是：找到树中所有处于指定层级（level）的基准节点，为每个基准节点生成一条从根到该节点的分支，分支需保留基准节点的所有子树，同时修剪掉路径上其他无关分支；若分支未达指定深度，则保留该分支的完整结构。

解法思路 #

1. 收集目标路径：递归遍历树，收集所有到达指定层级的基准节点的完整路径（从根到基准节点的节点链），同时收集未达指定层级的叶子节点路径。
2. 构建分支：对每个收集到的路径，生成对应的分支：
   • 路径上的非基准节点：只保留通往基准节点的子节点，修剪其他兄弟分支。
   • 基准节点/叶子节点：保留其所有子树（如果有）。
3. 返回结果：将所有生成的分支整理为列表返回。

伪代码实现 #

FUNCTION copy_tree(node):
    IF NOT IS_LIST(node):
        RETURN node
    RETURN [copy_tree(item) FOR item IN node]

FUNCTION build_branch(path):
    IF path IS EMPTY:
        RETURN EMPTY LIST
    current_node = path[0]
    IF LENGTH(path) == 1:
        # 基准节点或叶子节点，返回完整节点（值+所有子节点）
        RETURN [current_node[0]] + copy_tree(current_node[1:])
    ELSE:
        # 非最后一个节点，只保留路径中的下一个子节点
        next_node = path[1]
        child_branch = build_branch(path[1:])
        RETURN [current_node[0]] + [child_branch]

FUNCTION split_tree(tree, target_level):
    result = EMPTY LIST
    path_so_far = EMPTY LIST

    FUNCTION collect_target_paths(current_node, current_depth):
        ADD current_node TO path_so_far

        IF current_depth == target_level:
            branch = build_branch(path_so_far)
            ADD branch TO result
            REMOVE LAST ITEM FROM path_so_far
            RETURN

        # 当前节点是叶子节点（无子女）
        IF LENGTH(current_node) == 1:
            branch = build_branch(path_so_far)
            ADD branch TO result
            REMOVE LAST ITEM FROM path_so_far
            RETURN

        # 遍历所有子节点继续递归
        FOR child IN current_node[1:]:
            collect_target_paths(child, current_depth + 1)

        REMOVE LAST ITEM FROM path_so_far

    # 特殊处理层级为0的情况：直接返回原树副本
    IF target_level == 0:
        RETURN [copy_tree(tree)]

    collect_target_paths(tree, 0)
    RETURN result

Python 实现 #

def copy_tree(node):
    """递归复制嵌套列表表示的树，避免修改原树"""
    if not isinstance(node, list):
        return node
    return [copy_tree(item) for item in node]

def build_branch(path):
    """根据节点路径构建目标分支"""
    if not path:
        return []
    current_node = path[0]
    if len(path) == 1:
        # 基准节点或叶子节点，保留所有子节点
        return [current_node[0]] + copy_tree(current_node[1:])
    else:
        # 非基准节点，仅保留通往目标节点的子分支
        child_branch = build_branch(path[1:])
        return [current_node[0]] + [child_branch]

def split_tree(tree, target_level):
    result = []
    path_so_far = []

    def collect_target_paths(current_node, current_depth):
        path_so_far.append(current_node)

        # 到达目标层级，生成分支
        if current_depth == target_level:
            branch = build_branch(path_so_far)
            result.append(branch)
            path_so_far.pop()
            return

        # 当前节点是叶子节点，且未达目标层级，生成完整分支
        if len(current_node) == 1:
            branch = build_branch(path_so_far)
            result.append(branch)
            path_so_far.pop()
            return

        # 遍历所有子节点继续递归
        for child in current_node[1:]:
            collect_target_paths(child, current_depth + 1)

        # 回溯，移除当前节点
        path_so_far.pop()

    # 处理根节点为基准节点的情况
    if target_level == 0:
        return [copy_tree(tree)]

    collect_target_paths(tree, 0)
    return result

# 测试示例
if __name__ == "__main__":
    # 定义示例树结构
    tree = [
        "ROOT",
        ["B", ["E"], ["F"]],
        ["C"],
        ["D", ["H"], ["I", ["J"], ["K"]]]
    ]

    print("=== 测试 level=1 ===")
    for idx, branch in enumerate(split_tree(tree, 1), 1):
        print(f"分支{idx}: {branch}")

    print("\n=== 测试 level=2 ===")
    for idx, branch in enumerate(split_tree(tree, 2), 1):
        print(f"分支{idx}: {branch}")

测试输出 #

=== 测试 level=1 ===
分支1: ['ROOT', ['B', ['E'], ['F']]]
分支2: ['ROOT', ['C']]
分支3: ['ROOT', ['D', ['H'], ['I', ['J'], ['K']]]]

=== 测试 level=2 ===
分支1: ['ROOT', ['B', ['E']]]
分支2: ['ROOT', ['B', ['F']]]
分支3: ['ROOT', ['C']]
分支4: ['ROOT', ['D', ['H']]]
分支5: ['ROOT', ['D', ['I', ['J'], ['K']]]]

说明 #

• 树的结构要求：每个节点为嵌套列表，格式为[节点值, 子节点1, 子节点2, ...]，叶子节点格式为[节点值]。
• 函数返回的所有分支都是原树的副本，不会修改输入的原始树。
• 若指定层级大于树的最大深度，会返回所有叶子节点对应的完整分支。

内容的提问来源于stack exchange，提问作者Justcurious