对具有多个输入的函数使用df.apply生成多个输出

2024-10-01 13:40:53 发布

您现在位置:Python中文网/ 问答频道 /正文

我有一个像这样的数据框

initial year0 year1
0       0     12
1       1     13
2       2     14
3       3     15

请注意,year列的数目year0、year1。。。(year_count)是完全可变的,但在整个代码中是不变的

我首先想对每个“year”列应用一个函数来生成“mod”列,如下所示

def mod(year, scalar):
    return (year * scalar)

s = 5
year_count = 2
# Generate new columns
df[[f"mod{y}" for y in range (year_count)]] = df[[f"year{y}" for y in range(year_count)]].apply(mod, scalar=s)  

initial year0 year1 mod0 mod1
0       0     12    0   60
1       1     13    5   65
2       2     14    10  70
3       3     15    15  75

到目前为止一切都很好。问题是我现在想对year列和它对应的mod列应用另一个函数来生成另一组val列,比如

def sum_and_scale(year_col, mod_col, scale):
    return (year_col + mod_col) * scale

然后我将其应用于每个列(year0,mod0),(year1,mod1)等,以生成下一批列

当比例=10时,我应该以

initial year0 year1 mod0 mod1 val0 val1
0       0     12    0    60   0    720
1       1     13    5    65   60   780
2       2     14    10   70   120  840
3       3     15    15   75   180  900

这就是我遇到的问题——我不知道如何将两个现有的df列放在一个与第一个示例中的结构相同的函数中,如果我这样做的话

df[['val0', 'val1']] = df['col1', 'col2'].apply(lambda x: sum_and_scale('mod0', 'mod1', scale=10))

我不知道如何将其推广到具有任意输入和输出,以及应用常量比例参数。(我知道最后一条不起作用,但这是我见过的解决方案的另一条途径)

我问这个问题的原因是因为我相信我目前正在进行的循环会造成列数和每列长度的性能问题

谢谢


Tags: 函数moddfreturndefcountcolyear