Цикл по многомерному массиву и фильтрация по условию

Question

Я считаю, что похожие вопросы были заданы, но ни один из них не касается проблемы, с которой я сталкиваюсь.

У меня есть массив формы (H,W,L) - я должен пройтись по каждому экземпляру массива, чтобы отфильтровать значения (их координаты x, y), которые соответствуют определенным критериям. (скажем, val > t_r и val < t_c) - я должен повторить это для каждого из значений K.

Например: если у нас есть массив формы (2,3,4)

A = [[[1,2,3,], [3,4,5,]], 
        [[6,7,8],[1,4,5]], 
        [[5,7,7],[9,4,3]],
        [[1,2,4],[4,6,7]]]

предположим, что первый критерий равен val > 2, а второй критерий - val < 6 и сохраните значение (строка, столбец) в массиве N x 3. Где первые 2 значения 'row','col', а последнее соответствует слою / третьему измерению. тогда ожидаемый результат операции должен быть примерно таким -

output = [[0,2,0],[1,0,0],[1,0,0],[1,1,0],[1,2,0]....] это будет соответствовать значениям, отфильтрованным из A[:,:,0]

Один из подходов, о которых я подумал, - использование 3 для циклов- i, j, k, чтобы зациклить каждый из элементов, но я не могу понять точную реализацию. Я также хотел бы реализовать векторизацию везде, где это возможно. Я мог бы использовать некоторые указания.

Answer 1

Вы можете использовать np.nonzero и векторизовать свои сравнения.

a = np.asarray(A)
res = np.vstack(np.nonzero((a>2)&(a<6))).T

---

array([[0, 0, 2],
       [0, 1, 0],
       [0, 1, 1],
       [0, 1, 2],
       [1, 1, 1],
       [1, 1, 2],
       [2, 0, 0],
       [2, 1, 1],
       [2, 1, 2],
       [3, 0, 2],
       [3, 1, 0]], dtype=int64)

Вы всегда можете изменить порядок столбцов по своему вкусу, например:

res[:, [1,2,0]]

---

array([[0, 2, 0],
       [1, 0, 0],
       [1, 1, 0],
       [1, 2, 0],
       [1, 1, 1],
       [1, 2, 1],
       [0, 0, 2],
       [1, 1, 2],
       [1, 2, 2],
       [0, 2, 3],
       [1, 0, 3]], dtype=int64)