Search code examples
pythonsortingdataframetop-n

Python dataframe find index of top-5, then index into another column


I have a dataframe with two numeric columns, A & B. I want to find the top 5 values from col A and return the values from Col B held in the location of those top 5.

Many thanks.


Solution

  • I think need DataFrame.nlargest with column A for top 5 rows and then select column B:

    df = pd.DataFrame({'A':[4,5,26,43,54,36,18,7,8,9],
                       'B':range(10)})
    
    print (df)
        A  B
    0   4  0
    1   5  1
    2  26  2
    3  43  3
    4  54  4
    5  36  5
    6  18  6
    7   7  7
    8   8  8
    9   9  9
    

    print (df.nlargest(5, 'A'))
        A  B
    4  54  4
    3  43  3
    5  36  5
    2  26  2
    6  18  6
    
    a = df.nlargest(5, 'A')['B']
    print (a)
    4    4
    3    3
    5    5
    2    2
    6    6
    Name: B, dtype: int64
    

    Alternative solution with sorting:

    a = df.sort_values('A', ascending=False)['B'].head(5)
    print (a)
    4    4
    3    3
    5    5
    2    2
    6    6
    Name: B, dtype: int64