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

Question

我有一个长度为N的1D 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)

这是一个简洁明了的解决方案，但在我的电脑上没有速度优势。

- monade

---

np.isclose / np.allclose？ - jdehesa

感谢您的建议。请查看我的编辑，尝试使用np.isclose来实现它。 - monade

另一个选项是对数组进行排序（O(n log n)），然后运行滑动窗口（O(n)），计算窗口内的元素。总体上它将是O(n log n)，相比于O(n^2)的解决方案，但无法有效地向量化，因此在15k数组上可能会更慢。 - Marat

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)

预分类进一步改进：

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)

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

- Divakar

---

这太棒了，加速效果惊人。谢谢！ - monade

对于未来的谷歌员工：显然，我们必须将“out”除以“N”，尽管从速度上来说这并不是很重要。 - monade