python - 使用列表理解替换 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