Как ускорить итерацию по Python numpy .ndarray с 2 измерениями

Question

Итак, я просто хочу сделать это быстрее:

for x in range(matrix.shape[0]):
        for y in range(matrix.shape[1]):
            if matrix[x][y] == 2 or matrix[x][y] == 3 or matrix[x][y] == 4 or matrix[x][y] == 5 or matrix[x][y] == 6:
                if x not in heights:
                    heights.append(x)

Просто переберите матрицу 2x2 (обычно вокруг 18x18 или 22x22) и проверьте, что это x. Но это довольно медленно, мне интересно, какой самый быстрый способ сделать это.

Большое спасибо!

Answer 1

Для подхода, основанного на numpy, вы можете сделать:

np.flatnonzero(((a>=2) & (a<=6)).any(1))
# array([1, 2, 6], dtype=int64)

---

Где:

a = np.random.randint(0,30,(7,7))

print(a)

array([[25, 27, 28, 21, 18,  7, 26],
       [ 2, 18, 21, 13, 27, 26,  2],
       [23, 27, 18,  7,  4,  6, 13],
       [25, 20, 19, 15,  8, 22,  0],
       [27, 23, 18, 22, 25, 17, 15],
       [19, 12, 12,  9, 29, 23, 21],
       [16, 27, 22, 23,  8,  3, 11]])

---

Синхронизация в большем массиве:

a = np.random.randint(0,30, (1000,1000))

%%timeit
heights=[]
for x in range(a.shape[0]):
        for y in range(a.shape[1]):
            if a[x][y] == 2 or a[x][y] == 3 or a[x][y] == 4 or a[x][y] == 5 or a[x][y] == 6:
                if x not in heights:
                    heights.append(x)
# 3.17 s ± 59.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%%timeit
yatu = np.flatnonzero(((a>=2) & (a<=6)).any(1))
# 965 µs ± 11.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

np.allclose(yatu, heights)
# true

Векторизация с numpy дает примерно 3200x ускорение

Answer 2

Похоже, вы хотите найти, если в матрице появятся 2, 3, 4, 5 или 6.

Вы можете использовать np.isin(), чтобы создать матрицу истинных / ложных значений, а затем использовать это в качестве индексатора:

>>> arr = np.array([1,2,3,4,4,0]).reshape(2,3)
>>> arr[np.isin(arr, [2,3,4,5,6])]
array([2, 3, 4, 4])

При желании превратить это в простой Python set() для более быстрого in поиска и без дубликатов.

Чтобы получить позиции в массиве, где эти цифры появляются, используйте argwhere:

>>> np.argwhere(np.isin(arr, [2,3,4,5,6]))
array([[0, 1],
       [0, 2],
       [1, 0],
       [1, 1]])