You need to enable JavaScript to run this app.
最新活动
大模型
产品
解决方案
定价
生态与合作
支持与服务
开发者
了解我们

MultiIndex DataFrame下各公司列间运算并新增列的Pythonic实现

解决Pandas MultiIndex DataFrame高效添加公司专属计算列的问题

嘿,我来帮你搞定这个MultiIndex的扩展需求!针对你提到的公司数量多、数据量大的场景,我给你推荐两种高效且Pythonic的解决方案,完美适配你的需求,还能避免你之前遇到的拼接MultiIndex的麻烦。

方法一:直接通过MultiIndex元组赋值(最简洁高效)

这种方法利用Pandas对MultiIndex列的原生支持,直接通过元组索引给每个公司添加新列,全程是矢量化操作,速度快、内存占用低,非常适合上千家公司的场景:

import pandas as pd
import numpy as np

# 先构造你的示例数据
cols = pd.MultiIndex.from_product([['Company A','Company B'],['VWAL','Volumn']],names=[u'Entity',u'Indicator'])
rows = pd.date_range(start='2018-01-01',periods=6,freq='D')
df = pd.DataFrame(np.random.randint(1,100,(6,4)),index=rows,columns=cols)

# 获取所有唯一的公司名称
companies = df.columns.get_level_values('Entity').unique()

# 循环给每个公司添加计算列
for comp in companies:
    # 直接用MultiIndex元组作为列名,Pandas会自动归类到对应公司的层级下
    df[(comp, 'VWAL+Volumn')] = df[(comp, 'VWAL')] + df[(comp, 'Volumn')]
    df[(comp, 'VWAL-Volumn')] = df[(comp, 'VWAL')] - df[(comp, 'Volumn')]

# 排序列索引,让每个公司的所有列集中在一起(可选,但更整洁)
df = df.sort_index(axis=1)

执行后,你的DataFrame会自动扩展为每个公司包含4列的结构,示例输出如下:

Entity        Company A                     Company B                   
Indicator        VWAL Volumn VWAL+Volumn VWAL-Volumn        VWAL Volumn VWAL+Volumn VWAL-Volumn
2018-01-01          92      3          95          89          22     59          81         -37
2018-01-02          90     67         157          23          52     69         121         -17
2018-01-03          12     10          22           2          35     11          46          24
2018-01-04          83      7          90          76          62      5          67          57
2018-01-05          35     74         109         -39          27     19          46           8
2018-01-06          97     50         147          47          93     39         132          54

方法二:生成器迭代+分组计算(内存友好型)

如果你的数据量极端庞大,内存吃紧,可以用生成器来逐个生成新列,避免一次性把所有新列加载到内存中:

def generate_calculated_cols(df):
    # 按公司层级分组遍历
    for entity, group in df.groupby(level='Entity', axis=1):
        # 计算求和列,并设置对应的MultiIndex
        sum_col = group.xs('VWAL', level='Indicator', axis=1) + group.xs('Volumn', level='Indicator', axis=1)
        sum_col.columns = pd.MultiIndex.from_tuples([(entity, 'VWAL+Volumn')])
        yield sum_col
        
        # 计算差值列,同样设置MultiIndex
        diff_col = group.xs('VWAL', level='Indicator', axis=1) - group.xs('Volumn', level='Indicator', axis=1)
        diff_col.columns = pd.MultiIndex.from_tuples([(entity, 'VWAL-Volumn')])
        yield diff_col

# 合并原DataFrame和生成器产生的新列
df_expanded = pd.concat([df] + list(generate_calculated_cols(df)), axis=1).sort_index(axis=1)

生成器的优势在于按需生成数据,不会在内存中存储所有新列的副本,适合内存资源有限的场景。

为什么你的原代码遇到问题?

你之前尝试的temp = df.loc(axis=1)[:,'VWAL'].values+...得到的是一个numpy数组,丢失了原有的MultiIndex列信息,所以pd.concat时无法自动对应到各个公司的层级下。而上面的两种方法都保留了新列的MultiIndex结构,拼接时自然就能和原DataFrame的层级对齐。

内容的提问来源于stack exchange,提问作者Yan Tian

火山引擎 最新活动