Python 使用np.select选择两列

我试图解决熊猫数据帧的问题

我有一个数据框，它包含三列：

import numpy as np
np.random.seed(0)
dataframe = pd.DataFrame({'operation': ['data_a', 'data_b', 'avg', 'concat', 'sum', 'data_a', 'concat'], 
             'data_a': list(np.random.uniform(-1,1,[7,2])), 'data_b': list(np.random.uniform(-1,1,[7,2]))})

列“operation”表示合并列，因此，如果列“operation”中有“data_a”值，则表示取该特定行的data_a值，如果有“avg”操作，则取该特定行的“data_a”和“data_b”的平均值，依此类推

正如我在输出中所期望的，一个新的列包含与操作列的合并函数相同的值

我所尝试的：

dataframe['new_column'] = 'dummy_values'

for i in range(len(dataframe)):
    
    if dataframe['operation'].iloc[i]  == 'data_a':
        dataframe['new_column'].iloc[i] = dataframe['data_a'].iloc[i]
    elif dataframe['operation'].iloc[i] == 'data_b':
        dataframe['new_column'].iloc[i] = dataframe['data_b'].iloc[i]
    elif dataframe['operation'].iloc[i] == 'avg':
        dataframe['new_column'].iloc[i] = dataframe[['data_a','data_b']].iloc[i].mean()
    elif dataframe['operation'].iloc[i] == 'sum':
        dataframe['new_column'].iloc[i] = dataframe[['data_a','data_b']].iloc[i].sum()
    elif dataframe['operation'].iloc[i] == 'concat':
        dataframe['new_column'].iloc[i] = np.concatenate([dataframe['data_a'].iloc[i], dataframe['data_b'].iloc[i]], axis=0)

上面的解决方案相当慢，所以我尝试了np.select方法，如下所示

import numpy as np
con1 = dataframe['operation']  == 'data_a'
con2 = dataframe['operation']  == 'data_b'


val1 = dataframe['data_a']
val2 = dataframe['data_b']

dataframe['new_column'] = np.select([con1,con2], [val1,val2])

但如果我使用np.select选择两列，则会出现错误：

import numpy as np
con1 = dataframe['operation']  == 'data_a'
con2 = dataframe['operation']  == 'data_b'
con3 = dataframe['operation']  == 'avg'



val1 = dataframe['data_a']
val2 = dataframe['data_b']
val3 = dataframe[['data_b', 'data_a']].mean()


dataframe['new_column'] = np.select([con1,con2,con3], [val1,val2,val3])

错误消息

ValueError: shape mismatch: objects cannot be broadcast to a single shape

---

如何使用np选择不同的条件。选择？

检查

轴=1

，确保所有条件和值的形状相同

import numpy as np
con1 = dataframe['operation']  == 'data_a'
con2 = dataframe['operation']  == 'data_b'
con3 = dataframe['operation']  == 'avg'



val1 = dataframe['data_a']
val2 = dataframe['data_b']
val3 = dataframe[['data_b', 'data_a']].mean(axis = 1)


dataframe['new_column'] = np.select([con1,con2,con3], [val1,val2,val3])

---

dataframe['data\u b'，'data\u a']].平均值（axis=1）？@BEN\u YO添加一个答案，我会接受。什么是

take\u a

？@MadPhysicator更正了。什么是

concat

？谢谢，Bdw在性能和时间复杂性方面有比这更好的解决方案吗？我正在处理相当大的N-dim阵列并寻找快速解决方案。谢谢，这是执行时间的巨大飞跃。：）我刚刚检查了函数的值，axis=1给出了NaN，请检查output@AadityaUra如果行包含所有NaN，则平均值为NaN