如何在numpy矩阵中有效地找到行的最左非空值的索引列?

时间:2018-08-22 15:39:43

标签: python numpy

我想知道是否有一种快速的方法来找到每行最左的非空值的索引,它是一个numpy矩阵。

可以说我有这个:

data = np.array([[np.nan, np.nan, 10,      1],
                 [np.nan, 30,      50,      np.nan],
                 [40,      np.nan, np.nan, 5],
                 [np.nan, np.nan, np.nan, 10]])

我想分别返回10、30、40和10的索引,即2、1、0和3。

以下方法有效,但是效率不高(例如fn函数内部的循环很慢):

def fn(row):
    for idx, v in enumerate(row):
         if not np.isnan(v):
              return idx
np.apply_along_axis(fn, 1, data)

1 个答案:

答案 0 :(得分:5)

使用np.isnan获取非空的布尔数组,然后使用argmax(1)获取每一行的第一个索引,最后索引到输入数组-

In [19]: idx = (~np.isnan(data)).argmax(1)

In [20]: idx
Out[20]: array([2, 1, 0, 3])

In [21]: data[np.arange(len(idx)), idx]
Out[21]: array([10., 30., 40., 10.])