Python 提高索引搜索和删除的性能

我有一个相当简单的代码块，我想提高它的性能。它由一个

for

块组成，用于查找数组中整数的索引

下面的代码可以工作，但我觉得使用

for

将元素添加到空列表中并不是最好的方法

此块由MCMC使用，因此执行了数百万次。一个小的改进会变成一个大的改进。这能提高效率吗

import numpy as np

N = 20
# Integers from 1 to N
ran_indexes = np.random.randint(1, N, 1000)
# Number of integers to remove
rm_number = np.random.randint(0, 100, N)

# Better performance for this block?
# For each integer from 1 to N, keep only 'd' indexes of 'ran_indexes' that
# contain that integer, where 'd' is the ith element in 'rm_number'
new_indexes = []
for i, d in enumerate(rm_number):
    new_indexes += list(np.where(ran_indexes == i + 1)[0][:d])

---

在我的测试运行中，为新的索引预先分配空间和边走边填充始终比在现有列表中添加20-30%要好，请参见下面的实现

def f1（）：
新的_索引=[]
对于枚举中的i，d（rm_编号）：
新索引+=列表（np.where（ran_索引==i+1）[0][：d]）
返回新的索引
def f2（）：
新索引=np.0（总和（rm_数））
ind=0
对于枚举中的i，d（rm_编号）：
tmp=np.where（ran_index==i+1）[0][：d]
新索引[ind:ind+tmp.shape[0]]=tmp
ind+=tmp.shape[0]
返回列表（新索引[0:ind]）
在[144]中：%timeit f1
每个回路33.5纳秒±1.71纳秒（7次运行的平均值±标准偏差，每个10000000个回路）
在[145]：%timeit f2中
每个回路23.6纳秒±0.273纳秒（7次运行的平均值±标准偏差，每个10000000个回路）
在[146]中：%timeit f1
每个回路35.2纳秒±3.74纳秒（7次运行的平均值±标准偏差，每个10000000个回路）
在[147]：%timeit f2中
每个回路24.5纳秒±1.47纳秒（7次运行的平均值±标准偏差，每个10000000个回路）

另一方面：for循环的最后一次迭代没有变化，因此，

rm_numbers

中的最后一个数字永远不会用于任何有效率的事情。ran_索引中的最大数量为19，在上一次迭代中，您将检查

ran_索引==19+1

，它将始终为零。我不确定这是否是有意的，我想您应该修改ran_索引的定义，将

N+1

作为上界（假设上界是排他性的）。

---

如果19确实应该是最高的随机数，那么您应该能够跳过最后一个循环，从而缩短几纳秒的时间

列表联接，因为它们每次都需要一个全新的列表。更常见的情况是，在迭代构建数组时，我们使用list append，它已经就位，并且每次只在列表中添加元素

In [45]:  
    ...: new_indexes = [] 
    ...: for i, d in enumerate(rm_number): 
    ...:     new_indexes.append(np.where(ran_indexes == i + 1)[0][:d]) 
    ...:                                                                        
In [46]: new_indexes                                                            
Out[46]: 
[array([  5,  96, 143, 150, 154, 175]),
 array([ 14,  22,  26,  28,  32,  38,  46,  54,  70, 205, 218, 242, 248,
        254, 271, 318, 344, 352, 357, 393, 419, 437, 448, 472, 473, 503,
        521, 548, 558, 629, 631, 654, 661, 685, 699, 743, 755]),
 array([ 24,  34,  72,  97, 120, 140, 173, 181, 193, 199, 200, 225, 239,
        251, 265, 296, 350, 386, 411, 422, 465, 476, 506, 533, 609, 628,
        680, 694, 713, 759]),
 ....

通过这种构造，每个数组（

其中

结果）的长度不同，上界为

rm_编号

：

In [89]: [len(i) for i in new_indexes]-rm_number                                
Out[89]: 
array([  0,   0,   0,   0,   0,   0,   0,  -2, -24, -40,   0,  -3, -40,
         0, -15,  -5,   0,   0,   0, -96])

像这样的可变长度数组/列表很好地表明，你不能做一个超快速的“矢量化”（整个数组）操作，至少如果没有显著的技巧是不行的

我们可以获得您的代码生成的平面列表：

In [50]: np.concatenate(new_indexes).shape                                      
Out[50]: (626,)

一些时间安排：

In [53]: %%timeit  
    ...: new_indexes = [] 
    ...: for i, d in enumerate(rm_number): 
    ...:     new_indexes += list(np.where(ran_indexes == i + 1)[0][:d]) 
    ...:                                                                        

320 µs ± 7.93 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [54]:                                                                        
In [54]: %%timeit 
    ...: new_indexes = [] 
    ...: for i, d in enumerate(rm_number): 
    ...:     new_indexes.append(np.where(ran_indexes == i + 1)[0][:d]) 
    ...:                                                                        

184 µs ± 268 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [55]:                                                                        
In [55]: %%timeit 
    ...: new_indexes = [] 
    ...: for i, d in enumerate(rm_number): 
    ...:     new_indexes.append(np.where(ran_indexes == i + 1)[0][:d]) 
    ...: new_indexes=np.concatenate(new_indexes) 
    ...:  
    ...:                                                                        
193 µs ± 622 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [79]: timeit f2()  # Lukas                                                          
291 µs ± 1.43 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

===

查找所有匹配项，并

np。其中（temp）[0]

是索引。但这不适用于您的

rm\u编号

界限

np.where(temp.T)[1]    # without the `rm_number` truncation  

np.where(temp[:,i])[0][:d]

---

我想到了两种可能的方法-删除for循环（

f_2

-30%更快）或使用numba（

f_3

-6倍更快）。Numba还需要一种稍微不同的实现方法——更少的python，更少的复制数据的工作，更多的numpy，更多的读取数据的工作。不确定你是否可以使用numba，但值得一试。当然，麻木的替代品是Cython。然而，Cython需要更多的重构，而不仅仅是用numba包装函数

import numba as nb
import numpy as np

def f_1(ran_indexes, rm_number):
    new_indexes = []
    for idx, qty in enumerate(rm_number):
        new_indexes += list(np.where(ran_indexes == idx + 1)[0][:qty])
    return new_indexes

def f_2(ran_indexes, rm_number):
    return np.hstack([np.where(ran_indexes == idx + 1)[0][:qty] for idx, qty in enumerate(rm_number)])

@nb.njit
def f_3(ran_indexes, rm_number):
    ans = np.zeros(rm_number.sum(), dtype=np.int64)
    count = 0
    for idx in range(rm_number.shape[0]):
        count_2 = 0
        for idx_2 in range(ran_indexes.shape[0]):
            if count_2 == rm_number[idx]:
                break
            if ran_indexes[idx_2] == idx + 1:
                ans[count + count_2] = idx_2
                count_2 += 1
        count += count_2
    return ans[:count]

if __name__ == '__main__':

    N = 20
    ran_indexes_ = np.random.randint(1, N, 1000)
    rm_number_ = np.random.randint(0, 100, N - 1)

    ans_1 = f_1(ran_indexes_, rm_number_)
    ans_2 = f_2(ran_indexes_, rm_number_)
    ans_3 = f_3(ran_indexes_, rm_number_)

    # check results
    print(sum(ans_1), sum(ans_2), sum(ans_3))
    print(len(ans_1), len(ans_2), len(ans_3))

结果:

%timeit f_1(ran_indexes_, rm_number_)
111 µs ± 279 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit f_2(ran_indexes_, rm_number_)
77 µs ± 118 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit f_3(ran_indexes_, rm_number_)
17 µs ± 6.01 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

---

选择此选项，因为它是最快的实现。在我的测试中，它比我的原始代码快约15%，而其余的都低于10%。谢谢大家!！

%timeit f_1(ran_indexes_, rm_number_)
111 µs ± 279 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit f_2(ran_indexes_, rm_number_)
77 µs ± 118 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit f_3(ran_indexes_, rm_number_)
17 µs ± 6.01 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)