使用np.select选择两个pandas列

Question

我正在尝试解决熊猫数据帧的问题，

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

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]))})

​

​

列' data_a‘表示合并列，因此如果列' 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.select选择不同的条件？

Answer 1

检查axis = 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])