Search code examples
pythonpandasindexingsubsetindices

Access entries in pandas data frame using a list of indices

I facing the issue that I need only a subset of a my original dataframe that is distributed over different rows and columns. E.g.:

# My Original dataframe
import pandas as pd
dfTest = pd.DataFrame([[1,2,3],[4,5,6],[7,8,9]])

Output:

   0  1  2
0  1  2  3
1  4  5  6
2  7  8  9

I can provide a list with rows and column indices where my desired values are located:

array_indices = [[0,2],[1,0],[2,1]]

My desired output is a series:

3
4
8

Can anyone help?

Solution

  • Use pd.DataFrame.lookup

    dfTest.lookup(*zip(*array_indices))

array([3, 4, 8])

    Which you can wrap in a pd.Series constructor

    pd.Series(dfTest.lookup(*zip(*array_indices)))

0    3
1    4
2    8
dtype: int64

    Slight variant

    i, j = np.array(array_indices).T
dfTest.values[i, j]

array([3, 4, 8])

    Similarly as above

    pd.Series(dfTest.values[i, j])

0    3
1    4
2    8
dtype: int64