具有可变列索引的数据帧的行总和 (Python)

Question

我有一个包含几百万行的数据框。我需要计算从特定列索引到最后一列的每一行的总和。每行的列索引是唯一的。具有所需输出的示例如下：

import pandas as pd

df = pd.DataFrame({'col1': [1, 2, 2, 5, None, 4],
                   'col2': [4, 2, 4, 2, None, 1],
                   'col3': [6, 3, 8, 6, None, 4],
                   'col4': [9, 8, 9, 3, None, 5],
                   'col5': [1, 3, 0, 1, None, 7],
                   })

df_ind = pd.DataFrame({'ind': [1, 0, 3, 4, 3, 5]})

for i in df.index.to_list():
    df.loc[i, "total"] = df.loc[i][(df_ind.loc[i, "ind"]).astype(int):].sum()

print(df)

>>
   col1  col2  col3  col4  col5  total
0   1.0   4.0   6.0   9.0   1.0   20.0
1   2.0   2.0   3.0   8.0   3.0   18.0
2   2.0   4.0   8.0   9.0   0.0    9.0
3   5.0   2.0   6.0   3.0   1.0    1.0
4   NaN   NaN   NaN   NaN   NaN    0.0
5   4.0   1.0   4.0   5.0   7.0    0.0

如何在不使用 for 循环的情况下使用 pandas 有效地实现此目的。谢谢

Answer 1

您可以创建一个列出所有列位置的类似索引的 DataFrame，然后通过将此 DataFrame 与 df_ind 进行比较，您可以为整个原始 DataFrame 创建一个掩码。

然后 mask 原始 DataFrame 和 sum 根据因行而异的适当索引位置获取行总和。

import pandas as pd

mask = pd.DataFrame({col: df.columns.get_loc(col) for col in df.columns}, 
                    index=df.index)
#   col1  col2  col3  col4  col5
#0     0     1     2     3     4
#1     0     1     2     3     4
#2     0     1     2     3     4
#3     0     1     2     3     4
#4     0     1     2     3     4
#5     0     1     2     3     4

mask = mask.ge(df_ind['ind'], axis=0)
#    col1   col2   col3   col4   col5
#0  False   True   True   True   True
#1   True   True   True   True   True
#2  False  False  False   True   True
#3  False  False  False  False   True
#4  False  False  False   True   True
#5  False  False  False  False  False

df['total'] = df[mask].sum(1)

print(df)
   col1  col2  col3  col4  col5  total
0   1.0   4.0   6.0   9.0   1.0   20.0
1   2.0   2.0   3.0   8.0   3.0   18.0
2   2.0   4.0   8.0   9.0   0.0    9.0
3   5.0   2.0   6.0   3.0   1.0    1.0
4   NaN   NaN   NaN   NaN   NaN    0.0
5   4.0   1.0   4.0   5.0   7.0    0.0

Answer 2

另一种选择，使用 numpy：


cols = np.arange(len(df.columns))
# build a 2D array
mask = np.tile(cols, (len(df), 1))
# generate booleans by comparing to `df_ind`
mask = mask >= df_ind.to_numpy()
# replace True with `df`
mask = np.where(mask, df, mask)
# convert nulls to zero, and sum along the columns
mask = np.nan_to_num(mask).sum(1)
df.assign(total = mask)

   col1  col2  col3  col4  col5  total
0   1.0   4.0   6.0   9.0   1.0   20.0
1   2.0   2.0   3.0   8.0   3.0   18.0
2   2.0   4.0   8.0   9.0   0.0    9.0
3   5.0   2.0   6.0   3.0   1.0    1.0
4   NaN   NaN   NaN   NaN   NaN    0.0
5   4.0   1.0   4.0   5.0   7.0    0.0

具有可变列索引的数据帧的行总和 (Python)

Row sums of dataframe with variable column indexes (Python)

python

row

sum

pandas