Python实现树形数据结构的层级聚合函数应用

现有一种树形数据结构，节点包含整数标量、列表、字典及numpy数组。需要编写Python函数，将聚合函数应用于树形结构的底层节点，并用聚合结果替换底层节点，返回新的数据结构。示例如下：

# 初始数据结构
{"a": [1, [2, 3, 4], [5, 6, 7]], "b": [{"c":8, "d":9}, {"e":3, "f":4}, 8]}
# 聚合函数：sum

第一次调用结果:  {"a": [1, 9, 18], "b": [17, 7, 8]}
第二次调用结果:  {"a": 28, "b": 32}
第三次调用结果:  60
第四次调用结果:  60

解决方案 #

通过递归遍历树形结构实现核心逻辑：判断当前节点是否为仅含标量的底层节点，若是则直接返回；若不是则递归处理子节点，待子节点全部处理完成后，若所有子节点均为标量则应用聚合函数合并，否则保留处理后的结构。

代码实现 #

import numpy as np

def aggregate_tree(data, agg_func):
    # 处理标量节点
    if isinstance(data, (int, float)):
        return data
    # 处理numpy数组（直接视为标量集合进行聚合）
    elif isinstance(data, np.ndarray):
        return agg_func(data)
    # 处理列表：递归处理每个元素后判断是否全为标量
    elif isinstance(data, list):
        processed_items = [aggregate_tree(item, agg_func) for item in data]
        if all(isinstance(item, (int, float)) for item in processed_items):
            return agg_func(processed_items)
        return processed_items
    # 处理字典：递归处理每个值后判断是否全为标量
    elif isinstance(data, dict):
        processed_dict = {k: aggregate_tree(v, agg_func) for k, v in data.items()}
        if all(isinstance(val, (int, float)) for val in processed_dict.values()):
            return agg_func(processed_dict.values())
        return processed_dict
    # 其他未定义类型直接返回
    else:
        return data

测试示例 #

# 初始数据
initial_data = {"a": [1, [2, 3, 4], [5, 6, 7]], "b": [{"c":8, "d":9}, {"e":3, "f":4}, 8]}

# 第一次调用
result1 = aggregate_tree(initial_data, sum)
print("第一次调用结果:", result1)  # 输出: {"a": [1, 9, 18], "b": [17, 7, 8]}

# 第二次调用
result2 = aggregate_tree(result1, sum)
print("第二次调用结果:", result2)  # 输出: {"a": 28, "b": 32}

# 第三次调用
result3 = aggregate_tree(result2, sum)
print("第三次调用结果:", result3)  # 输出: 60

# 第四次调用
result4 = aggregate_tree(result3, sum)
print("第四次调用结果:", result4)  # 输出: 60

代码说明 #

• 底层节点判定：直接识别整数、浮点数；numpy数组直接作为标量集合处理。
• 递归逻辑：列表和字典分别遍历子元素/值进行递归处理，确保每一层的底层节点都被聚合。
• 聚合触发条件：仅当某一层的所有子节点处理后均为标量时，才会对该层应用聚合函数，否则保留原结构层级。

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