Search code examples
pythonpandasdataframerowindices

Find integer row-index from pandas index

The following code find index where df['A'] == 1

import pandas as pd
import numpy as np
import random

index = range(10)
random.shuffle(index)
df = pd.DataFrame(np.zeros((10,1)).astype(int), columns = ['A'], index = index)

df.A.iloc[3:6] = 1
df.A.iloc[6:] = 2

print df

print df.loc[df['A'] == 1].index.tolist()

It returns pandas index correctly. How do I get the integer index ([3,4,5]) instead using pandas API?

   A
8  0
4  0
6  0
3  1
7  1
1  1
5  2
0  2
2  2
9  2
[3, 7, 1]
Solution

  • Here is one way:

    df.reset_index().index[df.A == 1].tolist()

    This re-indexes the data frame with [0, 1, 2, ...], then extracts the integer index values based on the boolean mask df.A == 1.

    Edit Credits to @Max for the index[df.A == 1] idea.