pandas中基于树状层级规则迭代计算Result结果列的实现方法

核心思路 #

• 树结构的节点值计算遵循从下到上的顺序：最底层叶子节点没有子节点，Result直接等于自身Value；上层父节点的Result依赖直接子节点的Result计算结果，所以优先计算层级最高的节点，再逐步向上计算更低层级的父节点。
• 直接子节点判定规则：对于层级为n的父节点，其直接子节点层级一定为n+1，且子节点前n个LEVEL列的取值和父节点对应LEVEL列的取值完全一致。

具体代码实现 #

首先导入依赖并构造示例数据：

import pandas as pd
import numpy as np

# 构造示例输入数据
data = {
    'LEVEL': [1,2,3,4,5,6,6,5,6,6],
    'LEVEL-1': ['A01', np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
    'LEVEL-2': [np.nan, 'A011', np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
    'LEVEL-3': [np.nan, np.nan, 'AO111', np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
    'LEVEL-4': [np.nan, np.nan, np.nan, 'A01111', np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
    'LEVEL-5': [np.nan, np.nan, np.nan, np.nan, 'AO11111', np.nan, np.nan, 'AO11112', np.nan, np.nan],
    'LEVEL-6': [np.nan, np.nan, np.nan, np.nan, np.nan, 'AO111111', 'AO111112', np.nan, 'AO111121', 'AO111122'],
    'Value': [10,20,30,40,50,60,70,80,90,100]
}
df = pd.DataFrame(data)

然后执行计算逻辑：

# 初始化Result列
df['Result'] = np.nan

# 获取所有层级，从最高层级开始计算
max_level = df['LEVEL'].max()
# 最高层级是叶子节点，直接赋值Result=Value
df.loc[df['LEVEL'] == max_level, 'Result'] = df.loc[df['LEVEL'] == max_level, 'Value']

# 从次高层级逐级向下计算到最低层级（LEVEL=1）
for level in range(max_level-1, 0, -1):
    # 遍历当前层级的所有节点
    for idx, row in df[df['LEVEL'] == level].iterrows():
        # 筛选当前节点的直接子节点：层级为level+1，且前level个LEVEL列取值和父节点完全一致
        mask = (df['LEVEL'] == level + 1)
        for l in range(1, level+1):
            col = f'LEVEL-{l}'
            mask &= (df[col] == row[col])
        # 计算当前节点Result = 自身Value + 所有直接子节点Result之和
        child_result_sum = df.loc[mask, 'Result'].sum()
        df.loc[idx, 'Result'] = row['Value'] + child_result_sum

结果验证 #

计算后的Result列和题目给出的预期结果完全一致：

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