[Python] What is the fastest way to find the group by max in a column in a Python Pandas...

UPDATE2: I actually have 2000 draws not 3.

UPDATE: My df column A was wrong. I fixed it.

I have a really large version of df below.

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


I am trying to find which of the draw columns wins for each draw. Below is my current attempt, but its slow, but works. I feel like there should be a way with apply or something to make it faster.

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


Desired Output:

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


I generate the 2000 columns like so:

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)

Continue reading...