Efficient subsetting of pandas dataframe with indexes in 2d numpy array

user_7473630 0
Viewed 140

I have a 2d (in future it may be a 3d) array with indicies. Let's say it looks like that:

[[1,1,1],
[1,2,2],
[2,2,3]]

And I have a pandas dataframe:

index, A, B
1, 0.1, 0.01
2, 0.2, 0.02
3, 0,3, 0.03

I want to get a numpy array (or pandas df) with values from column A, sliced based on numpy array. So the result would be here:

[[0.1,0.1,0.1],
[0.1,0.2,0.2],
[0.2,0.2,0.3]]

I can do it with a loop to get pandas dataframe:

pd.DataFrame(df.A[val].values for val in array)

However I'm looking for more efficient way to do it. Is there better way that allows me to use whole array of indices at once?

pandas dataframe numpy arrays python
1 Answers

You can do:

df.loc[a.ravel(),'A'].values.reshape(a.shape)

Output:

array([[0.1, 0.1, 0.1],
       [0.1, 0.2, 0.2],
       [0.2, 0.2, 0.3]])
user_42384080
Related