применение определенной функции к различным диапазонам строк в пандах

Question

У меня есть df, как это с тоннами строк:

    BB      AA    FF        
     2      5      0    
     3      7      A
     6      5      A
     9      6      A
     8      3      0

И такая функция:

def test(a,b):
    # a=array col AA
    # b=array col BB
    return (a*b)+a

Я хотел бы, чтобы для строк в столбце FF, где значенияare! = 0, чтобы применить функциональный тест к тому срезу (массиву) df, который включает столбец BB и AA, чтобы сгенерировать следующий вывод в новом столбце ZZ:

    BB      AA    FF   ZZ      
     2      5      0   0 
     3      7      A   28
     6      5      A   35
     9      6      A   51
     8      3      0   0

Я думал что-то вроде:

df['zz']= df.apply(lambda x: test(df.AA,df.BB) for the range of values among zero)

Но моя проблема в том, что я не уверен, как указать де-массив в столбце FF для применения столбца

Answer 1

Вы можете использовать DataFrame.apply + маска :

def test(x):
    return (x[0]*x[1])+x[0]
df['ZZ']=df[['AA','BB']].apply(test,axis=1).mask(df['FF'].eq('0'),0)
print(df)

---

   BB  AA FF  ZZ
0   2   5  0   0
1   3   7  A  28
2   6   5  A  35
3   9   6  A  60
4   8   3  0   0

или вы можете использовать лямбда-функция :

df['ZZ']=df.apply(lambda x: x[['BB','AA']].prod()+ x['AA'] if x['FF'] != '0' else x['FF'],axis=1)
print(df)

---

   BB  AA FF  ZZ
0   2   5  0   0
1   3   7  A  28
2   6   5  A  35
3   9   6  A  60
4   8   3  0   0