Python 计算一个numpy数组中有多少个元素在每个其他元素的增量内

考虑数组

x

和增量变量

d

np.random.seed([3,1415])
x = np.random.randint(100, size=10)
d = 10

对于

x

中的每个元素，我想计算每个元素中有多少其他元素在距离delta

d

的范围内

所以x看起来像

print(x)

[11 98 74 90 15 55 13 11 13 26]

结果应该是

[5 2 1 2 5 1 5 5 5 1]

---

我所尝试的
战略：

• 使用广播来获取外部差异
• 外差绝对值
• 总共有多少超出了阈值

---

（np.abs（x[：，无）-x）策略

pir1(x, d)

[5 2 1 2 5 1 5 5 5 1]    


由于
x
不一定要排序，我们将通过
argsort
对其进行排序并跟踪排序排列，以便反转排列
我们将在
x
上使用
np.searchsorted
和
x-d
查找
x
值开始超过
x-d
的起始位置
在另一边再做一次，除非我们必须使用
np.searchsorted
参数
side='right'
并使用
x+d
取右搜索排序和左搜索排序之间的差值，计算每个元素+/-d范围内的元素数
使用argsort反转排序排列


将出现问题的方法定义为
pir1

def pir1(a, d):
    return (np.abs(a[:, None] - a) <= d).sum(-1)


演示

pir1(x, d)

[5 2 1 2 5 1 5 5 5 1]    



定时

pir2
显然是赢家

代码

pir1(x, d)

[5 2 1 2 5 1 5 5 5 1]    

功能



扩展到更大的阵列

摆脱了
pir1

from timeit import timeit

results = pd.DataFrame(
    index=np.arange(1, 11) * 1000,
    columns=['pir%s' %i for i in range(2, 5)])

for i in results.index:
    np.random.seed([3,1415])
    x = np.random.randint(1000000, size=i)
    for j in results.columns:
        setup = 'from __main__ import x, {}'.format(j)
        results.loc[i, j] = timeit('{}(x, 10)'.format(j), setup=setup, number=100)

results.insert(0, 'pir1', 0)

results.plot()

本文中列出了另外两个基于
搜索排序策略的变体

更有效的方法是从中获得
s.argsort（）
的有效思想

运行时测试-

In [155]: # Inputs
     ...: a = np.random.randint(0,1000000,(10000))
     ...: d = 10


In [156]: %timeit pir2(a,d) #@ piRSquared's post solution
     ...: %timeit pir3(a,d)
     ...: %timeit pir4(a,d)
     ...: 
100 loops, best of 3: 2.43 ms per loop
100 loops, best of 3: 4.44 ms per loop
1000 loops, best of 3: 1.66 ms per loop

非常感谢。我已经更新了测试结果以包括这些变体，并扩展了数组的大小。@piRSquared这些都是有意义的。对于较小的数组，
pir4
中的开销用于创建范围数组以获取
s.argsort（）
使它比简单的排序更不值钱。对于这个计数问题，您可以考虑使用
searchsorted
！
from timeit import timeit

results = pd.DataFrame(
    index=np.arange(1, 50),
    columns=['pir%s' %i for i in range(1, 5)])

for i in results.index:
    np.random.seed([3,1415])
    x = np.random.randint(1000000, size=i)
    for j in results.columns:
        setup = 'from __main__ import x, {}'.format(j)
        results.loc[i, j] = timeit('{}(x, 10)'.format(j), setup=setup, number=10000)

results.plot()

from timeit import timeit

results = pd.DataFrame(
    index=np.arange(1, 11) * 1000,
    columns=['pir%s' %i for i in range(2, 5)])

for i in results.index:
    np.random.seed([3,1415])
    x = np.random.randint(1000000, size=i)
    for j in results.columns:
        setup = 'from __main__ import x, {}'.format(j)
        results.loc[i, j] = timeit('{}(x, 10)'.format(j), setup=setup, number=100)

results.insert(0, 'pir1', 0)

results.plot()

def pir3(a,d):  # Short & less efficient
    sidx = a.argsort()
    p1 = a.searchsorted(a+d,'right',sorter=sidx)
    p2 = a.searchsorted(a-d,sorter=sidx)
    return p1 - p2

def pir4(a, d):   # Long & more efficient
    s = a.argsort()

    y = np.empty(s.size,dtype=np.int64)
    y[s] = np.arange(s.size)

    a_ = a[s]
    return (
        a_.searchsorted(a_ + d, 'right')
        - a_.searchsorted(a_ - d)
    )[y]

In [155]: # Inputs
     ...: a = np.random.randint(0,1000000,(10000))
     ...: d = 10


In [156]: %timeit pir2(a,d) #@ piRSquared's post solution
     ...: %timeit pir3(a,d)
     ...: %timeit pir4(a,d)
     ...: 
100 loops, best of 3: 2.43 ms per loop
100 loops, best of 3: 4.44 ms per loop
1000 loops, best of 3: 1.66 ms per loop