Есть ли способ использовать np.where () со списком столбцов?

Question

Я пытаюсь создать новую переменную с помощью функции np.where ().

myDF['newVar'] = np.where((myDF['var1']==1) |
                           (myDF['var2']==1) |
                           (myDF['var3']==1) ,
                           1, 0)

Есть ли способ заменить var1, var2, var3 списком, подобным этому (то же условие == 1 для каждого столбца) :

listVars=['var1', 'var2', 'var3']
myDF['newVar'] = np.where((myDF[listVars]==1),
                           1, 0)

Answer 1

Вы можете использовать pandas.DataFrame.any() с axis=1:

listVars=['var1', 'var2', 'var3']
myDF['newVar'] = np.where((myDF[listVars]==1).any(axis=1), 1, 0)

Например:

myDF = pd.DataFrame({
    "var1": [1, 1, 1, 0, 0],
    "var2": [1, 0, 1, 0, 1],
    "var3": [0, 0, 0, 0, 0]
})

listVars=['var1', 'var2', 'var3']
myDF['newVar'] = np.where((myDF[listVars]==1).any(1),1, 0)

print(myDF)
#   var1  var2  var3  newVar
#0     1     1     0       1
#1     1     0     0       1
#2     1     1     0       1
#3     0     0     0       0
#4     0     1     0       1