Delete numpy axis 1 based on condition-CodePudding

I need to remove values from a np axis based on a condition.

For example, I would want to remove [:,2] (the second values on axis 1) if the first value == 0, else I would want to remove [:,3].

Input:

[[0,1,2,3],[0,2,3,4],[1,3,4,5]]

Output:

[[0,1,3],[0,2,4],[1,3,4]]

So now my output has one less value on the 1st axis, depending on if it met the condition or not.

I know I can isolate and manipulate this based on array[np.where(array[:,0] == 0)] but then I would have to deal with each condition separately, and it's very important for me to preserve the order of this array.

I am dealing with 3D arrays & am hoping to be able to calculate all this simultaneously while preserving the order.

Any help is much appreciated!

CodePudding user response：

A possible solution:

a = np.array([[0,1,2,3],[0,2,3,4],[1,3,4,5]])
b = np.arange(a.shape[1])
np.apply_along_axis(
    lambda x: x[np.where(x[0] == 0, np.delete(b,2), np.delete(b,3))], 1, a)

Output:

array([[0, 1, 3],
       [0, 2, 4],
       [1, 3, 4]])

CodePudding user response：

Since you are starting and ending with a list, a straight forward iteration is a good solution:

In [261]: alist =[[0,1,2,3],[0,2,3,4],[1,3,4,5]]
In [262]: for row in alist:
     ...:     if row[0]==0: row.pop(2)
     ...:     else: row.pop(3)
     ...:     
In [263]: alist
Out[263]: [[0, 1, 3], [0, 2, 4], [1, 3, 4]]

A possible array approach:

In [273]: arr = np.array([[0,1,2,3],[0,2,3,4],[1,3,4,5]])

In [274]: mask = np.ones(arr.shape, bool)
In [275]: mask[np.arange(3),np.where(arr[:,0]==0,2,3)]=False

In [276]: mask
Out[276]: 
array([[ True,  True, False,  True],
       [ True,  True, False,  True],
       [ True,  True,  True, False]])

arr[mask] will be 1d, but since we are deleting the same number of elements each row, we can reshape it:

In [277]: arr[mask].reshape(arr.shape[0],-1)
Out[277]: 
array([[0, 1, 3],
       [0, 2, 4],
       [1, 3, 4]])

I expect the list approach will be faster for small cases, but the array should scale better. I don't know where the trade off is.