首页 > 解决方案 > 在 Python Pandas 数据框中的列中按 max 查找分组并标记它的最快方法是什么?

问题描述

UPDATE2:我实际上有 2000 次平局而不是 3 次。

更新:我的 df 列 A 是错误的。我修好了它。

我有一个非常大的版本df

data = {'A':[11111, 11111, 33333,11111], 'B':[101, 101, 102, 101],'C':[1,2,3,4],
    'draw0':[5, 6, 2, 1], 'draw1':[4,3,2,1], 'draw2':[2,3,4,6]}
df = pd.DataFrame(data)

     A     B   C  draw0   draw1   draw2
0  11111  101  1      5      4      2
1  11111  101  2      6      3      3
2  33333  102  3      2      2      4
3  11111  101  4      1      1      6

我试图找出每次平局中哪个平局列获胜。以下是我目前的尝试,但速度很慢,但有效。我觉得应该有一种方法可以应用或让它更快。

draw_cols = [col for col in df if col.startswith('draw')]

for col in draw_cols:
    max_idx = df.groupby(['A', 'B'])[col].idxmax().values
    df.loc[max_idx, col] = 1
    df.loc[~df.index.isin(max_idx), col] = 0

期望的输出:

     A     B   C  draw0  draw1  draw2
0  11111  101  1      0      1      0
1  11111  101  2      1      0      0
2  33333  102  3      1      1      1
3  11111  101  4      0      0      1

我像这样生成 2000 列:

def simulateDraw(df, n=2000):
    
    #simulate n drawings from the alpha and beta values and create columns 
    return pd.concat([df,
           df.apply(lambda row: pd.Series(np.random.beta(row.C, row.C, size=n)), axis = 1).add_prefix('draw')],
          axis = 1)

标签: pythonpandasperformance

解决方案


# groupby and transform the idxmax
max_idx = df.groupby(['A', 'B'])[df.columns[3:]].transform('idxmax')
# create a new column that is just your index
# this is done just in case your real data does not have a range index
max_idx['index'] = max_idx.index.values
# where the max_idx is in the index to return bool values and then update the original df
df.update(max_idx.isin(max_idx['index']).astype(int))

       A    B  C  draw0  draw1  draw2
0  11111  101  1      0      1      0
1  11111  101  2      1      0      0
2  33333  102  3      1      1      1
3  11111  101  4      0      0      1

推荐阅读