2
\$\begingroup\$

I have a DataFrame like this:

n=5
df_test = pd.DataFrame(index=range(n),data={'A': np.random.randn(n),
                                            'B': np.random.randn(n),
                                            'C': np.random.randn(n)})
cols = ['A','B','C']

It will looks something like this:

            A          B            C
0   -1.218522    0.458914   -1.299367
1    1.872055    0.202141   -0.441987
2    1.104548    1.513785   -2.106882
3   -0.511132   -0.888288   -0.740854
4    1.628356    0.191733    0.394444

I want to cycle through the columns and each time extract the maximum value until there are no rows left.

I have written the following code which works fine:

vals = np.zeros(n)
cols = []
idx = []
for row_idx in range(n):
    col_idx = row_idx%3
    selected_col = cols[col_idx]
    selected_idx = df_test.drop(idx)[selected_col].idxmax()
    idx.append(selected_idx)
    cols.append(selected_agent)
    vals[row_idx] = df_test.loc[selected_idx, selected_col]

print(idx)
print(cols)
print(vals)

Correctly returning:

[1, 2, 4, 3, 0]
['A', 'B', 'C', 'A', 'B']
[ 1.87205485  1.51378518  0.39444431 -0.5111323   0.45891356]

I can't help thinking that there is a neater way to do this. Especially growing those lists in the loop.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.