Python多节点树形结构组件总需求量快速计数算法咨询

算法思路 #

• 你的数据量仅100条，完全不存在性能压力，最优算法时间复杂度为O(n)，每个节点仅遍历1次，无需担心迭代次数过多的问题
• 核心逻辑：
  1. 先做两次O(n)的预处理，构建两个映射表：id映射表（存储每个组件id的基础属性）、父子映射表（存储每个父组件对应的所有子组件id列表）
  2. 从根节点出发用广度优先搜索（BFS）遍历整棵树，遍历过程中直接计算每个节点的总需求量，父节点计算完成后再处理子节点，确保子节点计算时能直接拿到父节点的总需求量

Python 实现代码 #

from collections import deque
import json

# 示例数据，已修正原始输入的语法错误
raw_data = [
    {"id": 123, "needed": 1, "need_id": None},
    {"id": 100, "needed": 2, "need_id": 123},
    {"id": 101, "needed": 3, "need_id": 123},
    {"id": 105, "needed": 3, "need_id": 123},
    {"id": 210, "needed": 5, "need_id": 101},
    {"id": 999, "needed": 2, "need_id": 210},
]

def calc_component_total(data):
    # 1. 构建id映射表，新增total字段存储总需求量
    id_map = {}
    # 2. 构建父子映射表
    parent_to_children = {}
    root_id = None

    for item in data:
        c_id = item["id"]
        need_id = item["need_id"]
        id_map[c_id] = {**item, "total": 0}
        # 定位根节点
        if need_id is None:
            root_id = c_id
            continue
        # 填充父子映射关系
        if need_id not in parent_to_children:
            parent_to_children[need_id] = []
        parent_to_children[need_id].append(c_id)
    
    # 3. BFS遍历计算总需求量
    q = deque()
    # 根节点总需求等于自身needed值
    id_map[root_id]["total"] = id_map[root_id]["needed"]
    q.append(root_id)

    while q:
        current_id = q.popleft()
        current_total = id_map[current_id]["total"]
        # 遍历所有子节点计算总需求
        for child_id in parent_to_children.get(current_id, []):
            child_node = id_map[child_id]
            child_node["total"] = current_total * child_node["needed"]
            q.append(child_id)
    
    # 返回格式：{组件id: 总需求量}
    return {c_id: node["total"] for c_id, node in id_map.items()}

# 测试运行
if __name__ == "__main__":
    # 如果是JSON文件，替换为：with open("your_file.json", "r", encoding="utf-8") as f: raw_data = json.load(f)
    result = calc_component_total(raw_data)
    print(result)
    # 输出结果：{123: 1, 100: 2, 101: 3, 105: 3, 210: 15, 999: 30} 与示例预期完全一致

补充说明 #

• 实现默认输入为合法树形结构，无循环依赖、无缺失父节点的异常数据，如果你的数据源存在异常可以自行添加校验逻辑，不会影响核心计算性能
• 若树深度极高，可将BFS替换为DFS递归实现，逻辑基本一致，仅遍历顺序不同

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