如何在每行中获取 n last/first 个真值的 NumPy 数组

Question

我有一个广播数组，它是按行排序的，还有一个屏蔽数组。我想获取每行的最后 n 个元素（或第一个 n 个元素） True 即：

a = np.array([[0.00298, 0.00455, 0.00767, 0.00939, 0.01104, 0.02351, 0.03370],
              [0.00298, 0.00455, 0.00767, 0.00939, 0.01104, 0.02351, 0.03370],
              [0.00298, 0.00455, 0.00767, 0.00939, 0.01104, 0.02351, 0.03370]])

mask = np.array([[1, 0, 0, 1, 1, 0, 1], [0, 1, 0, 0, 0, 1, 1], [0, 0, 1, 1, 0, 0, 0]], dtype=bool)

# a[mask] --> [0.00298  0.00939  0.01104  0.0337  0.00455  0.02351  0.0337  0.00767  0.00939]

# needed last  two --> [[0.01104  0.0337 ]  [0.02351  0.0337 ]  [0.00767  0.00939]]
# needed first two --> [[0.00298  0.00939]  [0.00455  0.02351]  [0.00767  0.00939]]

我们是否必须拆分数组（使用 np.cumsum(np.sum(mask, axis=1))）、填充和……？
使用 NumPy just 执行此操作的最佳方法是什么？

Answer 1

您可以通过将数组相乘并使用带有切片的列表理解来分别获取第一个或最后两个元素来实现：

# First two:
[[i for i in j if i][:2] for j in a*mask]
# Last two:
[[i for i in j if i][-2:] for j in a*mask]

输出：

[[0.00298, 0.00939], [0.00455, 0.02351], [0.00767, 0.00939]]
[[0.01104, 0.0337], [0.02351, 0.0337], [0.00767, 0.00939]]

编辑，仅使用 numpy 函数：

# First two:
np.apply_along_axis(lambda x: x[np.flatnonzero(x)[:2]], axis=1, arr=mask*a)
# Last two:
np.apply_along_axis(lambda x: x[np.flatnonzero(x)[-2:]], axis=1, arr=mask*a)

输出：

array([[0.00298, 0.00939], [0.00455, 0.02351], [0.00767, 0.00939]])
array([[0.01104, 0.0337 ], [0.02351, 0.0337 ], [0.00767, 0.00939]])

Answer 2

使用numpy获取前n个True：

n=2
a[(np.cumsum(mask, axis=1)<=n)&mask].reshape(-1,n)

输出：

array([[0.00298, 0.00939],
       [0.00455, 0.02351],
       [0.00767, 0.00939]])

最后一个：

n=2
a[(np.cumsum(mask[:,::-1], axis=1)<=n)[:,::-1]&mask].reshape(-1,n)

输出：

array([[0.01104, 0.0337 ],
       [0.02351, 0.0337 ],
       [0.00767, 0.00939]])

注意。每行必须至少有n个True才能得到正确的最终形状

如何在每行中获取 n last/first 个真值的 NumPy 数组

How to get NumPy array of n last/first Trues in each row

python

numpy