Итерации над выходом `np.where`

У меня есть 3D-массив и используйте np.where для поиска элементов, удовлетворяющих определенному условию. Вывод np.where представляет собой набор из трех 1D массивов, каждый из которых дает индексы вдоль одной оси. Я хотел бы перебрать этот вывод и распечатать индекс каждой точки в матрице, удовлетворяющей условию.

Один из способов сделать это:

indices = np.where(myarray == 0)
for i in range(0, len(indices[0])):
 print indices[0][i], indices[1][i], indices[2][i]

Однако, это выглядит немного громоздким, и мне было интересно, есть ли лучший способ?

2 ответа

Используйте zip

indices = zip(*np.where(myarray == 0))

Затем вы можете сделать

for i, j, k in indices:
 print ...

Например,

In [1]: x = np.random_integers(0, 1, (3, 3, 3))
In [2]: np.where(x) # you want np.where(x==0)
Out[2]: (array([0, 0, 0, 0, 0, 1, 1, 1, 1, 2]),
 array([0, 1, 1, 2, 2, 0, 0, 1, 1, 2]),
 array([1, 0, 1, 0, 1, 1, 2, 0, 2, 2]))
In [3]: zip(*np.where(x))
Out[3]: [(0, 0, 1),
 (0, 1, 0),
 (0, 1, 1),
 (0, 2, 0),
 (0, 2, 1),
 (1, 0, 1),
 (1, 0, 2),
 (1, 1, 0),
 (1, 1, 2),
 (2, 2, 2)]


Используйте np.transpose над zip, он быстрее для больших массивов

import numpy as np
myarray = np.random.randint(0, 7, size=1000000)
%timeit indices = zip(*np.where(myarray == 0))
%timeit indices = np.transpose(np.where(myarray == 0))
10 loops, best of 3: 31.8 ms per loop
100 loops, best of 3: 15.9 ms per loop

licensed under cc by-sa 3.0 with attribution.