Python Pandas dataframe：获取每列的最高 n 个非 NaN 值以及这些值的索引

Question

I have a pandas dataframe with values from multiple locations across many days.

import pandas as pd
import numpy as np
df = pd.DataFrame({'day': [1, 2, 3, 4, 5, 6],
                   'location-1': [10, 24, 24, 85, 90, np.NaN],
                   'location-2': [np.NaN, np.NaN, 45, 28, np.NaN, np.NaN]})
df.set_index('day', inplace=True)

I need to get the 4 highest values at each location , and the days on which they occur. NaN values need to be placed last. Something along the lines of:

result = pd.DataFrame({'location-1': [90, 85, 24, 24],
                       'location-2': [45, 29, np.NaN, np.NaN]})
result_days = pd.DataFrame({'location-1': [5, 4, 3, 2],
                            'location-2': [3, 4, 6, 5]})
    

I found a partial solution here: Get top 4 biggest values from each column using Pandas in Python

But that solution sorts NaN to the top and I can't find an na_position option for this. I saw solutions that then manually cycle each column's NaN down to the bottom, but I have no prior knowledge of which columns contain NaN and also have to keep track of the days. I can't use dropna because one location may have important values on the day that another location has NaN.

My questions are:

1. How do I sort this efficiently and extract the highest non-NaN values? I can hack it and replace NaN with -999 prior to sorting but I'd like to see if a general solution exists that doesn't rely on an assumption that my numbers are above a certain value.
2. How do I efficiently pull out the Days (or row indices) for the values in question 1? There may be repeat high values (as in location-1), and for that I need to go by latest day first. I have seen some solutions with np.argsort and np.argpartition, but I think they may hinge on how NaN are dealt with here.

This is my first time asking a question, and I am happy to clarify/change anything. Apologies if this is a duplicate; I hadn't found ones with the same situation. Thanks!

Answer 1

The following loop will give you what you need. You sort_values each location and assign it to the proper result and result_days

cols = ['location-1', 'location-2']
result = pd.DataFrame(columns=cols)
result_days  = pd.DataFrame(columns=cols)

for c in cols:
    tmp = df.sort_values(c, ascending=False).head(4)
    result[c] = tmp[c].values
    result_days[c] = tmp.index.values

print(result)
print(result_days)

   location-1  location-2
0        90.0        45.0
1        85.0        28.0
2        24.0         NaN
3        24.0         NaN
   location-1  location-2
0           5           3
1           4           4
2           2           1
3           3           2