CODEWITHSUNDEEP
pythonarrayslistnumpyconditional-statements
Tine
Tine

Reputation: 35

Get indices from array where element in row satisfies condition

I want to find the indexes of an array which satisfy a condition.

I have a numpy.ndarray B: (m = number of rows = 8 and 3 columns)

array([[ 0.,  0.,  0.],
   [ 0.,  0.,  0.],
   [ 0.,  0.,  1.],
   [ 0.,  1.,  1.],
   [ 0.,  1.,  0.],
   [ 1.,  1.,  0.],
   [ 1.,  1.,  1.],
   [ 1.,  0.,  1.],
   [ 1.,  0.,  0.]])

For each column I want to find the index of the rows for which the elements satisfy the following condition: for col in columns: B(row,col)=1 and B(row+1,col)=1 for all rows=1,2,..,m-1 and B(row,col)=1 for rows=0 and m.

So the desired outcome is:

Sets1=[[5, 6, 7, 8], [3, 4, 5], [2, 6]]

So far I've tried this:

Sets1=[]
for j in range(3):
    Sets1.append([i for i, x in enumerate(K[1:-1]) if B[x,j]==1 and B[x+1,j]==1])

But this is only the first part of the condition and gives the following wrong output because it takes the index of new set.. So it should be plus 1 actually..

Sets1= [[4, 5, 6], [2, 3, 4], [1, 5]]

Also the second part of the condition which goes for indexes 0 and m. Isn't included yet..

Edit: I fixed the plus 1 part by writing i+1 and I tried the second part of the condition by adding these if statements:

Sets1=[]
for j in range(3):
    Sets1.append([i+1 for i, xp in enumerate(K[1:-1]) if B[xp,j]==1 and B[xp+1,j]==1])
    if B[0,j]==1: Sets1[j].append(0)
    if B[(x-1),j]==1: Sets1[j].append(x-1)

Which does work since it gives the following output:

Sets1= [[5, 6, 7, 8], [3, 4, 5], [2, 6]]

So now i just need to add +1 to the elements of the list for the first part of the condition (before the if statements)...

I would really appreciate the help!

Upvotes: 2

Views: 2803

Answers (2)

user3483203
user3483203

Reputation: 51185

You can accomplish this with a boolean mask and np.where

First, the mask:

c1 = (x==1)
c2 = (np.roll(x, -1, axis=0) != 0)
c3 = (x[-1] == 1)

c1 & (c2 | c3)

array([[False, False, False],
       [False, False, False],
       [False, False,  True],
       [False,  True, False],
       [False,  True, False],
       [ True,  True, False],
       [ True, False,  True],
       [ True, False, False],
       [ True, False, False]])

np.where to get the indices:

>>> np.where(c1 & (c2 | c3))

(array([2, 3, 4, 5, 5, 6, 6, 7, 8], dtype=int64),
 array([2, 1, 1, 0, 1, 0, 2, 0, 0], dtype=int64))

If you really want the result as the list in your output:

s = np.where(c1 & (c2 | c3))
[list(s[0][s[1]==i]) for i in range(x.shape[1])]

# [[5, 6, 7, 8], [3, 4, 5], [2, 6]]

Upvotes: 1

Brenlla
Brenlla

Reputation: 1481

There is a vectorized way to do that with numpy. 1st we create a mask for where a equals 1:

mask=a.T==1.0

2nd mask will tell if next element also equals 1. Since we only want elements that satisfy both conditions we multiply both masks:

mask_next=np.ones_like(mask).astype(bool)
mask_next[:,:-1]=mask[:,1:]
fin_mask=mask*mask_next

Get the indexes:

idx=np.where(fin_mask)

1st index will tell us where to split the row idx:

split=np.where(np.diff(idx[0]))[0]+1
out=np.split(idx[1],split)

out produces the desired outcome. If I understood correctly, you want the indexes of elements equal to one where the next (columnwise) element is also one?

Upvotes: 1

Related Questions