advanced numpy/pandas indexing

Question

I would like to to perform an indexing operation something like

ix = [(1,2),(3,4),(5,6)]
ar[ix] # this is invalid real life

which would give me the 1-d array

array([ar[1,2], ar[3,4], ar[5,6]])

In other words I want to specify a set of coordinates and get back a vector of values at those coordinates. This is not work I am not fussy about the precise for of the index ix, any of list of pairs, pair of lists, 2-d array, pandas.DataFrame is fine. I am interested on doing this both on numpy arrays and Pandas DataFrames.

Answer 1

For numpy this can be accomplished in a very similar fashion:

>>> a = np.arange(5*5).reshape(5,5)
>>> ix = [(1,2),(3,4),(2,4)]
>>> a
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])

>>> a[zip(*ix)]
array([ 7, 19, 14])

For pandas you can acquire the values in a similar fashion by accessing the underlying numpy array:

>>> import pandas as pd
>>> df = pd.DataFrame(a)
>>> df.values[zip(*ix)]
array([ 7, 19, 14])

For a 2D array numpy expects a list of row indices and then column indices, with zip(*ix) you are providing this:

>>> zip(*ix)
[(1, 3, 2), (2, 4, 4)]