Numpy：将数组的每个元素与所有其他元素进行比较（±常数）

Question

我有一个长度为 N 的一维 Numpy 数组 A。对于数组中的每个元素x，我想知道数组中所有元素在[x-eps范围内的比例是多少； x+eps]，其中 eps 是常数。 N 大约为 15,000。

目前我是这样做的（最小的例子）：

import numpy as np

N = 15000
eps = 0.01
A = np.random.rand(N, 1)
prop = np.array([np.mean((A >= x - eps) & (A <= x + eps)) for x in A])

.. 在我的电脑上大约需要 1 秒。

我的问题：是否有更有效的方法？

编辑： 我认为 @jdehesa 在评论中的建议如下：

prop = np.isclose(A, A.T, atol=eps, rtol=0).mean(axis=1)

这是一个不错的简洁解决方案，但没有速度优势（在我的计算机上）。

Answer 1

这是利用 np.searchsorted -

的良好设置

sidx = A.argsort()
ridx = np.searchsorted(A, A+eps, 'right', sorter=sidx)
lidx = np.searchsorted(A, A-eps, 'left', sorter=sidx)
out = ridx - lidx

计时 -

In [71]: N = 15000
    ...: eps = 0.01
    ...: A = np.random.rand(N)

In [72]: %timeit np.array([np.sum((A >= x - eps) & (A <= x + eps)) for x in A])
560 ms ± 5.15 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [73]: %%timeit
    ...: sidx = A.argsort()
    ...: ridx = np.searchsorted(A, A+eps, 'right', sorter=sidx)
    ...: lidx = np.searchsorted(A, A-eps, 'left', sorter=sidx)
    ...: out = ridx - lidx
5.35 ms ± 47.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

进一步改进pre-sorting：

In [81]: %%timeit
    ...: sidx = A.argsort()
    ...: b = A[sidx]
    ...: ridx = np.searchsorted(b, A+eps, 'right')
    ...: lidx = np.searchsorted(b, A-eps, 'left')
    ...: out = ridx - lidx
3.93 ms ± 19.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

如评论中所述，对于 mean 等效版本，只需将最终数组输出除以 N。

Numpy：将数组的每个元素与所有其他元素进行比较（±常数）

Numpy: compare each element of array with all other elements (± constant)

python

arrays

performance

numpy

scientific-computing