获取符合我标准的 np.array 索引的最快方法答案

【问题标题】：Fastest way to get indices of a np.array that meet my criteria获取符合我标准的 np.array 索引的最快方法
【发布时间】：2020-06-05 22:55:09
【问题描述】：

我正在尝试找到最快的方法来获取满足我标准的矩阵元素的索引。我有一个 (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

因为我使用这个功能数万次，我需要它尽可能快地运行。

【问题讨论】：

标签： python performance numpy optimization indices

【解决方案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)

【讨论】：

当我运行您的代码时，仅将尺寸从 1000,1000 更改为 7,7。我的结果变为：每个循环 6.71 µs ± 94.6 ns（平均 ± 7 次运行的标准偏差，每个循环 100000 个循环）每个循环 12.8 µs ± 934 ns（平均 ± 7 次运行的标准偏差，每个循环 100000 个循环）

【解决方案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)

【讨论】：

查看更新时间。 np.int16 上的操作比 int64 上的慢。
我做了更多的测试，结果证明你是对的。我看到使用 int16 时提高了 ~19%，使用 int64 时提高了 ~21%。