首页 > 解决方案 > 如何将计算列添加到多索引数据框

问题描述

这可能已经在几个地方得到了回答,但是如果不求助于几行代码,我仍然无法让我的示例工作。

我有一个带有多索引列的数据框,我想向每个 0 级组添加一个计算列。

arrays = [
    np.array(["bar", "bar", "baz", "baz", "foo", "foo"]),
    np.array(["one", "two", "one", "two", "one", "two"]),
]

df = pd.DataFrame(np.linspace(0, 4*6-1, 4*6).reshape(4,6), columns=arrays)
print(df)

  bar     baz     foo    
  one two one two one two
0   0   1   2   3   4   5
1   6   7   8   9  10  11
2  12  13  14  15  16  17
3  18  19  20  21  22  23

具体来说,我想添加一列divisionbar,baz并按foo元素划分onetwo

我现在这样做的方式是这样的:

# divide the numbers of the two dataframes (have to drop the first level)
idx = pd.IndexSlice
div_res = df.loc[:, idx[:, "one"]].droplevel(1, axis=1) / df.loc[
    :, idx[:, "two"]
].droplevel(1, axis=1)
div_res.columns = pd.MultiIndex.from_tuples([(y, "division") for y in div_res.columns])

# and concatenate the result again
result = (
    pd.concat([df, div_res], axis=1).sort_index(
        level=[0, 1],
        axis=1,
    )
)
print(result)

       bar                baz                foo          
  division  one  two division  one  two division  one  two
0      0.0  0.0  1.0      0.7  2.0  3.0      0.8  4.0  5.0
1      0.9  6.0  7.0      0.9  8.0  9.0      0.9 10.0 11.0
2      0.9 12.0 13.0      0.9 14.0 15.0      0.9 16.0 17.0
3      0.9 18.0 19.0      1.0 20.0 21.0      1.0 22.0 23.0

这行得通,但在我看来,应该有一种更有效的方法来做到这一点。

非常感谢!

作为奖励,我还没有弄清楚如何对第二列级别进行排序:one, two, division.

标签: pythonpandasmulti-index

解决方案


您可以groupby在列轴上应用您的变​​换:

(df.groupby(level=0, axis=1) # groupby first level of columns
   .apply(lambda d: (d.xs('one', level=1, axis=1)        # select "one"
                      .div(d.xs('two', level=1, axis=1)) # divide by "two" 
                     ).rename(columns=lambda x: 'division')) # col name
   .join(df)  # join with original data
   .sort_index(level=0, axis=1)
)

输出:

        bar                   baz                   foo            
   division   one   two  division   one   two  division   one   two
0  0.000000   0.0   1.0  0.666667   2.0   3.0  0.800000   4.0   5.0
1  0.857143   6.0   7.0  0.888889   8.0   9.0  0.909091  10.0  11.0
2  0.923077  12.0  13.0  0.933333  14.0  15.0  0.941176  16.0  17.0
3  0.947368  18.0  19.0  0.952381  20.0  21.0  0.956522  22.0  23.0

推荐阅读