Pandas 多索引数据框列总和和分配

Question

考虑具有多索引列的数据框：

[ In 1]: df = pd.DataFrame(np.random.randint(0,10,[5,4]), 
                           columns=pd.MultiIndex.from_product([['A','B'], ['a','b']]))
         df

[Out 1]: 
            A       B
        a   b   a   b
    0   6   0   7   3
    1   0   5   1   8
    2   3   4   6   4
    3   6   2   0   2
    4   3   9   5   0

我想插入作为

level=1

列之和的新列。这是我的尝试，有效：

[ In 2]: df.loc[:, pd.IndexSlice['A','c']] = df.loc[:, pd.IndexSlice['A',:]].sum(axis=1)
         df.loc[:, pd.IndexSlice['B','c']] = df.loc[:, pd.IndexSlice['B',:]].sum(axis=1)
         df = df.sort_index(axis=1)
         df

[Out 2]: 
                A           B
        a   b   c   a   b   c
    0   6   0   6   7   3   10
    1   0   5   5   1   8   9
    2   3   4   7   6   4   10
    3   6   2   8   0   2   2
    4   3   9   12  5   0   5

是否有一种方法可以避免循环遍历

level=0

列名称？因为在我的真实数据框中，

level=0

有数百个文本列名称。

Python 3.8，pandas 1.4.4。

Answer 1

重塑数据框以便于计算

s = df.stack(0)
s['c'] = s.sum(axis=1)
s.stack().unstack([1, 2])

---

   A         B       
   a  b   c  a  b   c
0  2  9  11  1  9  10
1  2  2   4  8  5  13
2  4  8  12  9  7  16
3  9  1  10  8  5  13
4  1  0   1  5  8  13

Answer 2

您可以使用

groupby

操作并添加“c”级别。

out = df.combine_first(pd.concat({'c': df.groupby(level=0, axis=1).sum()}, axis=1)
                         .swaplevel(axis=1))

输出：

   A         B       
   a  b   c  a  b   c
0  2  0   2  6  5  11
1  0  8   8  5  0   5
2  6  5  11  3  4   7
3  4  3   7  2  1   3
4  7  2   9  5  9  14