使用列表推导替换 m x n 数组中的元素

Question

我有一个数组：

a = array([[1,2,3,1], [2,5,3,1], [0,0,0,0], [5,3,2,5]])

我想根据每行中的最后一项（0 或 1）遍历数组。如果每行的最后一项为0，我想将该行中的其他3项（仅）更改为np.nan。

例如：

a = array([[1,2,3,1], [2,5,3,1], [nan, nan, nan, 0], [5,3,2,5]])

我可以使用 for 循环来做到这一点。即：

for frames in range(len(a)):
    if a[frames][3] == 0:
        a[frames][0:2] = np.nan

有没有更有效的方法来使用列表推导来做到这一点？到目前为止，这就是我想出的全部，但觉得它可能会更有效率：

a = np.array([[np.nan, np.nan, np.nan] if frames[3] == 0 else frames[0:3] for frames in a])

因为这将创建一个数组数组，并裁剪最后一列

提前致谢！

Answer 1

IIUC，您不需要列表理解。使用indexing

>>> a = np.array([[1,2,3,1], [2,5,3,1], [0,0,0,0], [5,3,2,5]], dtype=float)
>>> a[a[:,-1] == 0, 0:3] = np.nan

array([[  1.,   2.,   3.,   1.],
       [  2.,   5.,   3.,   1.],
       [ nan,  nan,  nan,   0.],
       [  5.,   3.,   2.,   5.]])

---

如果您有这些数组的字典，只需索引它们中的每一个

for a in data.values():
    a[a[:,-1] == 0, 0:3] = np.nan