计算 numpy 数组中有多少元素在每个其他元素的 delta 范围内

Question

考虑数组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[:, None] - x) <= d).sum(-1)

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

---

这很好用。但是，它无法扩展。外部差异是 O(n^2) 时间。如何获得不随二次时间缩放的相同解决方案？

Answer 1

本文列出了另外两个变体，基于来自OP's answer post 的searchsorted strategy。

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]

更有效的方法派生出从this post 得到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

Answer 2

策略

• 由于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)

我们将定义一个新函数pir2

def pir2(a, d):
    s = x.argsort()
    a_ = a[s]
    return (
        a_.searchsorted(a_ + d, 'right')
        - a_.searchsorted(a_ - d)
    )[s.argsort()]

---

演示

pir1(x, d)

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

---

pir1(x, d)

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

---

时机
pir2 显然是赢家！

代码

功能

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

def pir2(a, d):
    s = x.argsort()
    a_ = a[s]
    return (
        a_.searchsorted(a_ + d, 'right')
        - a_.searchsorted(a_ - d)
    )[s.argsort()]

#######################
# From Divakar's post #
#######################
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]

---

测试

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()

---

扩展到更大的数组
摆脱了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()