Get top n columns and ignore NaN

Question

I am struggling to get the top n (n=3 in my case) columns while ignoring NaNs. My dataset:

import numpy as np
import pandas as pd
  
x = {'ID':['1','2','3','4','5'],
     'productA':[0.47, 0.65, 0.48, 0.58, 0.67],
     'productB':[0.65,0.47,0.55, np.NaN, np.NaN],
     'productC':[0.78, np.NaN, np.NaN, np.NaN, np.NaN],
     'productD':[np.NaN, np.NaN, 0.25, np.NaN, np.NaN],
     'productE':[0.12, np.NaN, 0.47, 0.12, np.NaN]}
       
df = pd.DataFrame(x)

My desired outcome:

ID	top3
A1	productC - productB - productA
A2	productA - productB
A3	productB - productA- productE
A4	productA - productE
A5	productA

As you can see, if n < 3, it should retain whatever n is but sorted by their values. I tried np.argsort but it doesn't ignore the NaNs and instead sorts the missing products by alphabetical order.

score 4 · Accepted Answer · edited Feb 17 '21 at 18:18

4

Try using :

df.set_index("ID").apply(
    lambda x: pd.Series(x.nlargest(3).index).tolist(), axis=1
)

ID
1    [productC, productB, productA]
2              [productA, productB]
3    [productB, productA, productE]
4              [productA, productE]
5                        [productA]
dtype: object

edited Feb 17 '21 at 18:18

Ch3steR

20,090
4
28
58

answered Feb 17 '21 at 17:58

Pygirl

12,969
5
30
43

1

that's a nice pandas solution, if performance is no issue (which it is unlikely to be) then it looks like a winner – Matt Feb 17 '21 at 18:29
@Matt: use swifter library. It will make the apply operation faster. – Pygirl Feb 17 '21 at 18:41

Ch3steR · Answer 2 · 2021-02-17T18:24:39.027

You can use np.argsort with np.isnan to filter out NaNs. Then just boolean indexing would do.

arr = df.iloc[:, 1:].to_numpy() # Leaving out `ID` col
idx = arr.argsort(axis=1)
m = np.isnan(arr)
m = m[np.arange(arr.shape[0])[:,None], idx]
out = df.columns[1:].to_numpy()[idx]
out = [v[~c][-3:] for v, c in zip(out, m)]

pd.Series(out, index= df['ID'])

ID
1    [productA, productB, productC]
2              [productB, productA]
3    [productE, productA, productB]
4              [productE, productA]
5                        [productA]
dtype: object

df.apply over axis=1 is just for-loop under the hood and can be slow. But you can leverage on NumPy functions(vectorized) to gain some efficiency.

In [152]: %%timeit 
     ...: df.set_index('ID').apply(lambda x: pd.Series(x.nlargest(3).index).toli
     ...: st(), axis=1) 
     ...:  
     ...:                                                                       
2.04 ms ± 19.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [153]: %%timeit 
     ...: arr = df.iloc[:, 1:].to_numpy() # Leaving out `ID` col 
     ...: idx = arr.argsort(axis=1) 
     ...: m = np.isnan(arr) 
     ...: m = m[np.arange(arr.shape[0])[:,None], idx] 
     ...: out = df.columns[1:].to_numpy()[idx] 
     ...: out = [v[~c][-3:] for v, c in zip(out, m)] 
     ...:  
     ...:                                                                       
144 µs ± 1.59 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Almost 14x performance gain.

Pandas apply with swifter library will give better performance. https://towardsdatascience.com/add-this-single-word-to-make-your-pandas-apply-faster-90ee2fffe9e8 :) — Pygirl, Feb 17 '21 at 18:43
Can you also add performance comparison with pandas apply + swifter ? — Pygirl, Feb 17 '21 at 18:44

score 2 · Answer 3 · answered Feb 17 '21 at 18:28

I would recommend just using numpy directly.

Depending on your experience you may find it a bit confusing and muddly (I certainly do)

import numpy as np

# your data
d = {
    'productA':[0.47, 0.65, 0.48, 0.58, 0.67],
    'productB':[0.65,0.47,0.55, np.NaN, np.NaN],
    'productC':[0.78, np.NaN, np.NaN, np.NaN, np.NaN],
    'productD':[np.NaN, np.NaN, 0.25, np.NaN, np.NaN],
    'productE':[0.12, np.NaN, 0.47, 0.12, np.NaN]
}

# replae your nans with -infs as otherwise they are counted as high
for k,v in d.items():
    d[k] = [-np.inf if i is np.NaN else i for i in v]

# store as a matrix
matrix = np.array(list(d.values()))

# your ids are 1 to 5
for i in range(1, 6):
    print(f"ID: {i}")
    
    # arg sort axis=0 will order how you want (by ooing over the horizontal axis)
    # you then want to select the i-1th column [::, i-1]
    # and do reverse order [::-1]
    print(np.argsort(matrix, axis=0)[::, i - 1][::-1])

Get top n columns and ignore NaN

3 Answers3