np.select 具有两个以上的 pandas 列

Question

我正在尝试解决熊猫问题陈述。熊猫的数据框是这样的：

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

列'operation'表示合并列，所以如果列'operation'中有'data_a'值，则表示取该特定行的data_a值，如果有'avg'操作，则取'data_a'的平均值以及该特定行的'data_b'，依此类推。

我对输出的期望，一个新列包含根据操作列的合并函数的值

我正在处理 NumPy 数组的第 n 个暗淡的许多行。

我尝试了两种解决方案，但都很慢。

第一个解决方案，用普通的python循环：

# first solution

start = time.time()
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)
        
end = time.time()
print(end - start)

# 0.3356964588165283

这很慢，第二种解决方案是pandas apply方法：

# second solution
start = time.time()
def f(x):
    if x['operation']  == 'data_a':
        return x['data_a']
    elif x['operation']  == 'data_b':
        return x['data_b']
    elif x['operation']  == 'avg':
        return x[['data_a','data_b']].mean()
    elif x['operation']  == 'sum':
        return x[['data_a','data_b']].sum()
    elif x['operation']  == 'concat':
        return  np.concatenate([x['data_a'], x['data_b']], axis=0)
        
dataframe['new_column'] = dataframe.apply(f, axis=1)

end = time.time()
print(end - start)

# 0.2401289939880371

这也很慢。我正在尝试使用 NumPy 选择方法来解决这个问题：

# third solution

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



val1 = dataframe['data_a']
val2 = dataframe['data_b']
val3 = dataframe[['data_b', 'data_a']].mean()
val4 = dataframe[['data_b', 'data_a']].sum()
val5 = dataframe[['data_b']]* dataframe[['data_a']]


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

这是错误的：

~/tfproject/tfenv/lib/python3.7/site-packages/numpy/lib/stride_tricks.py in _broadcast_shape(*args)
    189     # use the old-iterator because np.nditer does not handle size 0 arrays
    190     # consistently
--> 191     b = np.broadcast(*args[:32])
    192     # unfortunately, it cannot handle 32 or more arguments directly
    193     for pos in range(32, len(args), 31):

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

我该如何解决这个错误，有没有其他优化的方法来解决这个问题？

谢谢！

Answer 1

您可以使用 pandas 掩码对其进行矢量化，这样您只需执行所需的操作，但仍然具有矢量化的优势。为简洁起见，df 是您的数据框：

df['new_column'] = np.nan
mask = df['operation']=='data_a'
df.loc[mask, 'new_column'] = df.loc[mask, 'data_a']
mask = df['operation']=='data_b'
df.loc[mask, 'new_column'] = df.loc[mask, 'data_b']
mask = df['operation']=='avg'
df.loc[mask, 'new_column'] = (df.loc[mask, 'data_a'] + df.loc[mask, 'data_b'])/2
# etc