将函数的输出分配给与其他列中的条件匹配的行的数据框列

Question

我有一个数据框 'df'，其中包含 'A'、'B' 列，如图所示。我想创建一个新列 'C' 并将其分配给函数 func() 的输出。但是，限制是只有那些行（在 "C" 中）被分配这样的值，列 "A" 条目匹配 "abc"强>。对于此类行，func() 的输入变量将是 'B' 列的相应条目。

例如，在附加的照片中，第 0 行和第 2 行的列 'A' 值为 'abc'。对于此类行，'C' 中的值应由 func() 使用列 'B' 中的值计算（在此大小写，"2222" 或 "4444"）。

注意：变量 'var2'、'var3' 是在我的程序其他地方派生的一些值。我们可以在此处将它们视为常量进行讨论。

我当然可以使用“for 循环”来完成这项工作，但是，我想知道是否有任何有效的方法可以使用 pandas 或 numpy 来做到这一点。我尝试使用“.where”方法，但似乎 func() 将整个“B”列解释为其输入，而不是将“B”列的每个元素（即......“2222”或“4444”） .

def func(var1, var2, var3):
    <something>
    return(output)

import numpy as np
df['C'] = np.where(data['A'] == 'abc', func(data['B'], var2, var3), " ")

Answer 1

IIUC，您可以使用.loc 方法访问数据框中的每个项目，并使用.apply() 应用您的分配：

不知道你有多少条件，但模仿上面的代码：

df.loc[df['A'] == 'abc', 'C'] = df['B'].apply(func)

例如，

import pandas as pd
a = ['abc','b','abc']
b = [1,2,3]
df = pd.DataFrame({'A' : a,
              'B' : b})
print(df)
   A  B
0  abc  1
1    b  2
2  abc  3
df.loc[df['A'] == 'abc','C'] = df['B'].apply(lambda x : x ** 2)
   A    B    C
0  abc  1  1.0
1    b  2  NaN
2  abc  3  9.0

如果您的逻辑更复杂，您可以使用np.select，如下所示：

conditions  = [ df['B'] >= 3, (df['B'] < 3) & (df['B']> 1), df['B'] <= 1]
choices     = [ "high", 'medium', 'low' ]

df["C"] = np.select(conditions, choices, default=np.nan)

print(df)
     A  B       C
0  abc  1     low
1    b  2  medium
2  abc  3    high