根据索引列表获取df中列的值

Question

i have a list of indexes: [0,0,0,3,1,0,4,2,1] and a dataframe (this is just an example, the real dataframe is bigger...):

            snr        freq         snr  ...         snr        freq  freq_ref
0    111.796861  400.003168  116.805099  ...  123.952201  400.046262    400.00
1    111.800587  400.010109  117.194605  ...  124.033467  400.083761    400.05
2    111.636656  400.012101  117.654265  ...  124.155229  400.117228    400.10
3    111.839271  400.031985  118.009703  ...  124.208280  400.192227    400.15
4    112.162853  400.096895  118.196040  ...  124.055698  400.218755    400.20

i want to get the snr values according to the index list. each index is 1 snr values from a column: first index is the first snr column, 2nd index is the second snr column and the 3rd is the 3rd column. the 4th index is again the 1st snr column.

the output should be:

        snr           snr           snr        
0    111.796861    116.805099    123.952201
1    111.839271    117.194605    123.952201
2    112.162853    117.654265    124.033467

any ideas?

thanks

Answer 1

You can reshape list to 2d array first and then indexing in numpy - necessary 3 columns after selecting df['snr'] and number values in list is necessary len(a) == 3 * len(df.index) :

L =  [0,0,0,3,1,0,4,2,1]

a = np.array(L).reshape(-1,3)

df = pd.DataFrame(df['snr'].to_numpy()[a, np.arange(a.shape[1])])
print (df)
            0           1           2
0  111.796861  116.805099  123.952201
1  111.839271  117.194605  123.952201
2  112.162853  117.654265  124.033467