python - 对于每一行，找到第 n 个非 NaN 元素的列的最快方法是什么？

Question

我有一个 Python pandas DataFrame，其中每个元素都是 float 或 NaN。对于每一行，我需要找到包含该行第 n 个数字的列。也就是说，我需要获取包含非 NaN 的行的第 n 个元素的列。我知道第 n 个这样的列总是存在的。

因此，如果 n 为 4 并且名为 myDF 的 pandas 数据框如下:

      10   20   30   40   50   60  70  80  90  100

'A'  4.5  5.5  2.5  NaN  NaN  2.9 NaN NaN 1.1 1.8
'B'  4.7  4.1  NaN  NaN  NaN  2.0 1.2 NaN NaN NaN
'C'  NaN  NaN  NaN  NaN  NaN  1.9 9.2 NaN 4.4 2.1
'D'  1.1  2.2  3.5  3.4  4.5  NaN NaN NaN 1.9 5.5

我想获得:

'A'  60
'B'  70
'C'  100 
'D'  40

我能做到:

import pandas as pd
import math

n = some arbitrary int
for row in myDF.indexes:
   num_not_NaN = 0   
   for c in myDF.columns:    
      if math.isnan(myDF[c][row]) == False: 
           num_not_NaN +=1
      if num_not_NaN==n:
           print row, c
           break

我敢肯定这很慢而且不是很 Pythonic。如果我要处理非常大的 DataFrame 和很大的 n 值，有没有一种方法会更快？

Answer 1

如果速度是您的目标，那么尽可能使用 Pandas 的矢量化方法是个好主意:

>>> (df.notnull().cumsum(axis=1) == 4).idxmax(axis=1) # replace 4 with any number you like
'A'     60
'B'     70
'C'    100
'D'     40
dtype: object

其他答案都不错，而且在语法上可能更清晰一些。就速度而言，对于您的小例子，它们之间没有太大区别。然而，对于稍大的 DataFrame，矢量化方法已经快了大约 60 倍:

>>> df2 = pd.concat([df]*1000) # 4000 row DataFrame
>>> %timeit df2.apply(lambda row: get_nth(row, n), axis=1)
1 loops, best of 3: 749 ms per loop

>>> %timeit df2.T.apply(lambda x: x.dropna()[n-1:].index[0])
1 loops, best of 3: 673 ms per loop

>>> %timeit (df2.notnull().cumsum(1) == 4).idxmax(axis=1)
100 loops, best of 3: 10.5 ms per loop