基于特定条件的动态列的总和

Question

我试图根据某些条件获取动态列的总和。

cols = ['ID','A','B','C','D','E','F','G']

数据框df具有上面列出的所有列。

如果ID = 2，则需要前两列A，B的总和如果ID = 3，则需要前三列A，B，C的总和

df.loc['SUM'] = df.loc[df['ID'] > 0,cols[0:df['ID']]].sum(axis=1)

上面的代码行给出了以下错误：

TypeError: slice indices must be integers or None or have an __index__ method

注意： ID可以是任何数字，但始终小于或等于列总数

数据框中没有nan和缺少值。

Answer 1

借助 NumPy 的强大功能，您可以对其进行矢量化：

df.set_index('ID', inplace=True)
df.wherenp.arange(df.shape[1]) < df.index.values[:,None]).sum(axis=1)    

最小可重复示例

import pandas as pd

pd.__version__
# '0.24.2'

np.random.seed(0) 
df = pd.DataFrame(np.random.randint(1, 10, (5, 8))) 
df.insert(0, 'ID', np.random.randint(1, 8, 5))  
df                                           

   ID  0  1  2  3  4  5  6  7
0   4  6  1  4  4  8  4  6  3
1   7  5  8  7  9  9  2  7  8
2   4  8  9  2  6  9  5  4  1
3   3  4  6  1  3  4  9  2  4
4   5  4  4  8  1  2  1  5  8

df.set_index('ID', inplace=True)
df.assign(
    SUM=df.where(np.arange(df.shape[1]) < df.index.values[:,None]).sum(axis=1))

    0  1  2  3  4  5  6  7   SUM
ID                              
4   6  1  4  4  8  4  6  3  15.0
7   5  8  7  9  9  2  7  8  47.0
4   8  9  2  6  9  5  4  1  25.0
3   4  6  1  3  4  9  2  4  11.0
5   4  4  8  1  2  1  5  8  19.0

该where函数将隐藏我们不想求和的值：

df.where(np.arange(df.shape[1]) < df.index[:,None])   

    0  1  2    3    4    5    6   7
ID                                 
4   6  1  4  4.0  NaN  NaN  NaN NaN
7   5  8  7  9.0  9.0  2.0  7.0 NaN
4   8  9  2  6.0  NaN  NaN  NaN NaN
3   4  6  1  NaN  NaN  NaN  NaN NaN
5   4  4  8  1.0  2.0  NaN  NaN NaN

感谢@Quang Hoang 指出解决方案中的潜在错误。