获取满足我的条件的 np.array 索引的最快方法

Question

我正在尝试找到最快的方法来获取满足我的条件的矩阵元素的索引。我有一个 (7,7) np.array（名为 "board"），其中包含从 0 到 400 的 int16。例如，我想找到等于 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

在这种情况下，您似乎不需要索引，只需要一个掩码。

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)

Answer 2

如果您想最大程度地缩短函数的运行时间，最好的办法可能是避免在每次调用时分配新数组。这意味着在函数外部为临时值维护额外的数组，但它确实可以显着加快速度。

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)

获取满足我的条件的 np.array 索引的最快方法

Fastest way to get indices of a np.array that meet my criteria

python

optimization

performance

numpy

indices