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

I have a broadcasted array, which is sorted rowly, and a masked array. I want to get the last n elements (or the first n ones) of each row which are True i.e.:

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]]

Do we have to split the array (using np.cumsum(np.sum(mask, axis=1))), pad and …?
What will be the best way to do this just with NumPy?

CodePudding user response：

Using numpy to get the first n True:

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

Output:

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

Last n:

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

Output:

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

NB. There must be at least n True per row to have the correct final shape

CodePudding user response：

You can do it by multiplying the arrays and using a list comprehension with slicing to get the first or last two elements respectively:

# 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]

Output:

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

Edit, only using numpy functions:

# 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)

Output:

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]])