python - 对 NumPy 数组进行索引，忽略索引数组中的 NaN

Question

我有一个零数组

arr = np.zeros([5,5])
array([[ 0.,  0.,  0.,  0.,  0.],
   [ 0.,  0.,  0.,  0.,  0.],
   [ 0.,  0.,  0.,  0.,  0.],
   [ 0.,  0.,  0.,  0.,  0.],
   [ 0.,  0.,  0.,  0.,  0.]])

我想根据索引分配值，所以我这样做了。

out = np.array([[nan,2.,4.,1.,1.],[nan,3.,4.,4.,4.]])

arr[out[0].astype(int),np.arange(len(out[0]))] = 1
arr[out[1].astype(int),np.arange(len(out[1]))] = 1

如果有 0 而不是 nan，则赋值工作正常。

在 nan 的情况下如何跳过分配？是否可以立即从多维索引数组中赋值而不是使用 for 循环？

Answer 1

掩盖它 -

mask = ~np.isnan(out)
arr[out[0,mask[0]].astype(int),np.flatnonzero(mask[0])] = 1
arr[out[1,mask[1]].astype(int),np.flatnonzero(mask[1])] = 1

示例运行 -

In [171]: out
Out[171]: 
array([[ nan,   2.,   4.,   1.,   1.],
       [ nan,   3.,   4.,   4.,   4.]])

In [172]: mask = ~np.isnan(out)
     ...: arr[out[0,mask[0]].astype(int),np.flatnonzero(mask[0])] = 1
     ...: arr[out[1,mask[1]].astype(int),np.flatnonzero(mask[1])] = 1
     ...: 

In [173]: arr
Out[173]: 
array([[ 0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  1.,  1.],
       [ 0.,  1.,  0.,  0.,  0.],
       [ 0.,  1.,  0.,  0.,  0.],
       [ 0.,  0.,  1.,  1.,  1.]])

或者，用范围屏蔽替换 flatnonzero 调用 -

r = np.arange(arr.shape[1])
arr[out[0,mask[0]].astype(int),r[mask[0]]] = 1
arr[out[1,mask[1]].astype(int),r[mask[1]]] = 1

<小时/>

如果您处理的行数多于2，并且您想以矢量化方式分配它们，这里有一种方法，使用线性索引 -

n = arr.shape[1]
linear_idx = (out*n + np.arange(n))
np.put(arr, linear_idx[~np.isnan(linear_idx)].astype(int), 1)