加快在另一个索引列表中查找下一个索引（numpy）

Question

下面的代码可以正常工作，但由于循环，它似乎没有经过优化。我已经能够成功地矢量化我的所有其他方法，但我似乎无法弄清楚如何删除这个方法上的循环。

Speedwise：当我有数百万行时，这就成了一个问题。

有没有办法对此进行矢量化，或者我应该尝试使用 cython 或 numba？我一直在尝试限制使用的包数。

示例代码：

import numpy as np

leading = np.array([814, 935, 1057, 3069])

within = np.array([193, 207, 243, 251, 273, 286, 405, 427, 696, 770, 883,
                   896, 1004, 2014, 2032, 2033, 2046, 2066, 2079, 2154])

# find first following elements in within array
first_after_leading = []
for _ in leading:
    temp = (within - _).max()

    first_after_leading.append(temp)

# convert to np array
first_after_leading = np.array(first_after_leading)

Answer 1

leading中的每个元素对within中的所有元素减去的最大值将从within中的最大值减去leading。因此，只需执行 -

within.max() - leading

不需要额外的模块。

计时 -

In [79]: np.random.seed(0)
    ...: within = np.random.rand(1000000)
    ...: leading = np.random.rand(400000)

In [80]: %timeit within.max() - leading
1000 loops, best of 3: 850 µs per loop

Answer 2

我想把它做成单衬会有所帮助。试试吧。

first_after_leading =np.array([(within - _).max() for _ in leading])

Answer 3

使用 numba，您可以对代码进行相当简单的翻译：

import numba as nb
import numpy as np

def find_leading(leading, within):
    # find first following elements in within array
    first_after_leading = []
    for _ in leading:
        temp = (within - _).max()

        first_after_leading.append(temp)

    # convert to np array
    first_after_leading = np.array(first_after_leading)

    return first_after_leading


@nb.jit(nopython=True)
def find_leading_nb(leading, within):
    # find first following elements in within array
    first_after_leading = np.empty_like(leading)
    for i, _ in enumerate(leading):
        temp = (within - _).max()

        first_after_leading[i] = temp

    return first_after_leading

然后使用您的原始输入：

%timeit find_leading(leading, within)
%timeit find_leading_nb(leading, within)
%timeit (within[:,None] - leading).max(0)
17.3 µs ± 169 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
1.7 µs ± 25.3 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
6.48 µs ± 180 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

然后是一些更大的数组：

leading = np.random.randint(0, 100, (1000,))
within = np.random.randint(0, 100, (100000,))

%timeit find_leading(leading, within)
%timeit find_leading_nb(leading, within)
%timeit (within[:,None] - leading).max(0)
145 ms ± 3.82 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
67.4 ms ± 218 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
553 ms ± 4.42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

时间运行在 MacOS 上使用 numba 0.44 和 numpy 1.16.4 python 3.7

编辑

但是如果我正确理解你的算法，一个更快的方法是只找到 within 的最大值一次，然后用 leading 取差，所以你没有在循环中找到临时数组的max：

@nb.jit(nopython=True)
def find_leading_nb2(leading, within):
    max_within = within.max()
    first_after_leading = np.empty_like(leading)
    for i, x in enumerate(leading):
        first_after_leading[i] = max_within - x
    return first_after_leading

这会根据您的原始输入产生以下结果：

%timeit find_leading_nb2(leading, within)
919 ns ± 8.8 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

以及以下关于大输入的内容：

%timeit find_leading_nb2(leading, within)
21.6 µs ± 180 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

加快在另一个索引列表中查找下一个索引（numpy）

Speed up finding next index in another list of indices (numpy)

numpy

cython

python-3.x

numba