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)
# 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