Get values of pandas series from a array of index locations-CodePudding

I have a 2-d array of an index of a pandas series. Would like to create a 2-d array of the values from the pandas series that correspond to the index.

For example:

import pandas as pd
import numpy as np
A = pd.Series(data=[1,2,3,4,5])
idx = np.array([[0,2,3],[2,3,1]])

Would like to return:

B = np.array([[1,3,4],[3,4,2]])

I know I could do this as a loop:

B = np.zeros((2,3))
for i in [0,1]:
    B[i,:] = test[idx[i]]

However, in practice need to do this repeatedly so would like to broadcast the index locations directly. Pandas is not necessary, happy to do it all in numpy if easier.

CodePudding user response：

Something like this might work:

A[idx.flatten()].values.reshape(idx.shape)

CodePudding user response：

A[idx] gives a Cannot index with multidimensional key error.

In [190]: A = pd.Series(data=[1,2,3,4,5])
     ...: idx = np.array([[0,2,3],[2,3,1]])

But the 1d array derived from the Series, can be indexed this way:

In [191]: A.values
Out[191]: array([1, 2, 3, 4, 5])
In [192]: A.values[idx]
Out[192]: 
array([[1, 3, 4],
       [3, 4, 2]])

numpy has no problems returning an array with a dimension that matches idx.

Indexing the Series like this returns a Series - which by definition is 1d:

In [194]: A[idx.ravel()]
Out[194]: 
0    1
2    3
3    4
2    3
3    4
1    2
dtype: int64