Самый быстрый способ получить индексы np.array, которые соответствуют моим критериям

Question

Я пытаюсь найти самый быстрый способ получить индексы матричных элементов, которые соответствуют моим критериям. У меня есть (7,7) np.array (названный "board"), который содержит int16 от 0 до 400. Например, я хочу найти индексы элементов, которые равны 300.

Я перепробовал много техник, и до сих пор самый быстрый метод - это np.where (board == 300)

Функция, которую я пытаюсь оптимизировать:

def is_end(self, board):
    ind = np.where((board > 300) & (board - 300 < 100))
    try:
        victoriousPlayer = board[ind[0][0], ind[1][0]] % 100 // 10
        return victoriousPlayer
    except:
        return -1

Поскольку я использую эту функцию десятки тысяч раз мне нужно, чтобы он работал как можно быстрее.

Answer 1

Если вы хотите минимизировать время выполнения функции, вероятно, лучшее, что вы можете сделать, - это избегать выделения новых массивов при каждом вызове. Это означает, что поддержание дополнительных массивов для временных значений вне функции, но это дает вам значительное ускорение.

import numpy as np

# Original function
def is_end_1(board):
    ind = np.where((board > 300) & (board - 300 < 100))
    try:
        victoriousPlayer = board[ind[0][0], ind[1][0]] % 100 // 10
        return victoriousPlayer
    except:
        return -1

# Without array allocation
def is_end_2(board, tmpBool1, tmpBool2):
    np.less(300, board, out=tmpBool1)
    np.less(board, 400, out=tmpBool2)
    np.logical_and(tmpBool1, tmpBool2, out=tmpBool1)
    idx = np.unravel_index(np.argmax(tmpBool1), board.shape)
    return board[idx] % 100 // 10 if tmpBool1[idx] else -1

# Test
np.random.seed(0)
# Create some data
board = np.random.randint(500, size=(1000, 1000))
# Result from original function
res1 = is_end_1(board)
# Temporary arrays
tmpBool1 = np.empty_like(board, dtype=np.bool)
tmpBool2 = tmpBool1.copy()
# Result from function without allocations
res2 = is_end_2(board, tmpBool1, tmpBool2)
print(res1 == res2)
# True

# Measure time
%timeit is_end_1(board)
# 9.61 ms ± 323 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit is_end_2(board, tmpBool1, tmpBool2)
# 1.38 ms ± 53.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Answer 2

Вам не нужны индексы в этом случае, кажется, просто маска.

ind = np.where((board > 300) & (board - 300 < 100))
victoriousPlayer = board[ind[0][0], ind[1][0]] % 100 // 10

эквивалентно

victoriousPlayer = board[(board  > 300) & (board - 300 < 100)][0] % 100 // 10

Время:

In [1]: import numpy as np                                                                                                    

In [2]: board = np.random.randint(0,401, (7,7))                                                                               

In [3]: %timeit ind = np.where((board > 300) & (board - 300 < 100));victoriousPlayer = board[ind[0][0], ind[1][0]] % 100 // 10
6.77 µs ± 260 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [4]: %timeit victoriousPlayer = board[(board  > 300) & (board - 300 < 100)][0] % 100 // 10                                 
5.02 µs ± 26.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)