Python 为 EXCEL OFFSET 函数创建等效的 pandas 函数

发布于03月10日

假设输入是

d = {'col1': [1,2,3,4,5,6,7,8,9,10], 
 'col2': [1,2,3,4,5,6,7,8,9,10],
 'col3': [1,2,3,4,5,6,7,8,9,10],
 'offset': [1,2,3,1,2,3,1,2,3,1]} 

df = pd.DataFrame(data=d)

我想创建一个附加列，如下所示:

df['output'] = [1, 4, 9, 4, 10, 18, 7, 16, 27, 10]

基本上，offset中的每一个数字都告诉你要求和的列数(从col1作为参考点).

有没有一种矢量化的方法可以做到这一点，而不用遍历offset中的每个值？

推荐答案

你用np.select.要使用它，请创建sum列中的每个列(1、2、3…根据需要)作为可能的选项，并为偏移列中的每个值创建一个布尔掩码作为可能的条件.

# get all possible values from offset
lOffset = df['offset'].unique()

# get te result with np.select
df['output'] = np.select(
    # create mask for each values in offset
    condlist=[df['offset'].eq(i) for i in lOffset],
    # crerate the sum over the number of columns per offset value
    choicelist=[df.iloc[:,:i].sum(axis=1) for i in lOffset]
)
print(df)
#    col1  col2  col3  offset  output
# 0     1     1     1       1       1
# 1     2     2     2       2       4
# 2     3     3     3       3       9
# 3     4     4     4       1       4
# 4     5     5     5       2      10
# 5     6     6     6       3      18
# 6     7     7     7       1       7
# 7     8     8     8       2      16
# 8     9     9     9       3      27
# 9    10    10    10       1      10

注意:这假定偏移量列是最后一列