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

Question

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?

Quang Hoang · Accepted Answer · 2020-10-15 13:23:47Z

1

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

answered Oct 15, 2020 at 13:23

Quang Hoang

151k11 gold badges64 silver badges86 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Great, that's exactly what I was looking for. Thanks!