Получить индексы из массива, где элемент в строке удовлетворяет условию

Question

Я хочу найти индексы массива, которые удовлетворяют условию.

У меня есть numpy.ndarray B: (м = количество строк = 8 и 3 колонки)

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

Для каждого столбца я хочу найти индекс строк, для которых элементы удовлетворяют следующему условию: для столбца в столбцах: B (строка, столбец) = 1 и B (строка + 1, столбец) = 1 для всех строк = 1,2, .., m-1 и B (строка, столбец) = 1 для строк = 0 и m.

Итак, желаемый результат:

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

Пока я пробовал это:

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

Но это только первая часть условия и дает следующий неправильный вывод, потому что он принимает индекс нового набора .. Так что на самом деле это должно быть плюс 1 ..

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

Также вторая часть условия, которая идет для индексов 0 и m. Пока не включено ..

Редактировать: я исправил часть плюс 1, написав i + 1, и попробовал вторую часть условия, добавив следующие операторы if:

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)

Что работает, так как дает следующий вывод:

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

Так что теперь мне просто нужно добавить +1 к элементам списка для первой части условия (перед операторами if) ...

Буду очень признателен за помощь!

Answer 1

Вы можете сделать это с помощью логической маски и np.where

Сначала маска:

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 для получения индексов:

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

Если вы действительно хотите получить результат в виде списка в выводе:

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

Answer 2

Существует векторизованный способ сделать это с помощью numpy. Сначала мы создаем маску для которой a равно 1:

mask=a.T==1.0

2-я маска покажет, равен ли следующий элемент 1. Так как нам нужны только элементы, удовлетворяющие обоим условиям, мы умножаем обе маски:

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

Получить индексы:

idx=np.where(fin_mask)

1-й индекс скажет нам, где разбить строку idx:

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

out дает желаемый результат. Если я правильно понял, вы хотите, чтобы индексы элементов были равны единице, где следующий (по столбцу) элемент также равен единице?