Pandas分层多级索引：如何简洁实现跨行层级值的行计算？

问题描述 #

现有如下分层索引的DataFrame：

level1 level2
root   child1      10      10
       child2      20      20
       child3      30      30
       child4      40      40
root2  child1     100     100
       child2     200     200
       child3     300     300
       child4     400     400

需要新增desiredcol列，列值规则为：

• root/child1: child4的value1 + child1的value2
• root/child2: child1的value1 + child2的value2
• root/child3: child2的value1 + child3的value2
• root/child4: 0 + child4的value2
• root2下的子节点规则完全同上

目标行位置无规律，无法使用shift或rolling方法，已有可行方案但希望实现更简洁规范。

原实现代码 #

import pandas as pd
from io import StringIO

data = """
level1,level2,value1,value2
root,child1,10,10
root,child2,20,20
root,child3,30,30
root,child4,40,40
root2,child1,100,100
root2,child2,200,200
root2,child3,300,300
root2,child4,400,400
"""

df = pd.read_csv(StringIO(data), index_col=[0,1])

# We want the following result:
#level1,level2,value1,value2,desiredcol
#root,child1,10,10,(child4, value1) + (child1, value2)
#root,child2,10,10,(child1, value1) + (child2, value2)
#root,child3,10,10,(child2, value1) + (child3, value2)
#root,child4,10,10,(child4, value2)

# First - isolate the value1 column, and use unstack
# to flip the dimension columnwise. 
value1s = df[['value1']].unstack()
print('value1s unstacked:')
print(value1s)

# Copy the needed columns to the relevant mapped states.
value1s[('intermediate', 'child1')] = value1s[('value1', 'child4')]
value1s[('intermediate', 'child2')] = value1s[('value1', 'child1')]
value1s[('intermediate', 'child3')] = value1s[('value1', 'child2')]
value1s[('intermediate', 'child4')] = 0

# Get rid of the original values as they are still in df.
value1s = value1s.drop(columns=['value1'])
# Flip back to rows.
value1s = value1s.stack()
print(value1s)

df = df.join(value1s, how='inner')
# Now we can just row-wise sum the relevant columns.
df['desiredcol'] = df[['value2', 'intermediate']].sum(axis=1)
print(df)

优化后的简洁实现 #

利用手动映射字典直接关联目标行，避免不必要的维度转换操作，逻辑更直观：

import pandas as pd
from io import StringIO

data = """
level1,level2,value1,value2
root,child1,10,10
root,child2,20,20
root,child3,30,30
root,child4,40,40
root2,child1,100,100
root2,child2,200,200
root2,child3,300,300
root2,child4,400,400
"""

df = pd.read_csv(StringIO(data), index_col=[0,1])

# 定义每个level2对应的value1来源映射
source_map = {
    'child1': 'child4',
    'child2': 'child1',
    'child3': 'child2',
    'child4': None
}

# 计算desiredcol：根据映射获取对应value1（无来源则取0），与当前行value2求和
df['desiredcol'] = df.apply(
    lambda row: (df.loc[(row.name[0], source_map[row.name[1]]), 'value1'] 
                 if source_map[row.name[1]] is not None 
                 else 0) + row['value2'],
    axis=1
)

print(df)

优化优势 #

• 逻辑清晰：直接通过映射字典明确目标行关联关系，可读性强
• 代码简洁：减少中间变量和维度转换步骤，行数大幅减少
• 适配性强：完全支持目标行位置无规律的场景，手动映射规则一目了然

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