快速填充列表

Question

我有一小段代碼，我用它來填充整數列表。我需要提高它的性能，也許將整個事物轉化爲numpy陣列，但我不知道如何。

這裏的MWE：

import numpy as np 

# List filled with integers. 
a = np.random.randint(0,100,1000) 

N = 10 
b = [[] for _ in range(N-1)] 
for indx,integ in enumerate(a): 
    if 0<elem<N: 
     b[integ-1].append(indx) 

這是它做什麼：

• 在a
• 每一個整數（integ）看它是否位於一個給定的範圍之間(0,N)
• 如果是，則將其索引存儲在b的子列表中，其中索引所述子列表是用原來的整數減1（integ-1）

這段代碼運行非常快，但我實際的代碼使用更大的列表，從而改善其性能的需要。

Answer 1

試試這個解決方案！上半場我無恥地從喬的回答中偷走，但之後它使用排序和二進制搜索，這與N更好地擴展。

def new(a, N): 
    mask = (a > 0) & (a < N) 
    elements = a[mask] 
    indices = np.arange(a.size)[mask] 

    sorting_idx = np.argsort(elements, kind='mergesort') 
    ind_sorted = indices[sorting_idx] 

    x = np.searchsorted(elements, range(N), side='right', sorter=sorting_idx) 

    return [ind_sorted[x[i]:x[i+1]] for i in range(N-1)] 

你可以把x = x.tolist()在那裏額外雖然小速度的提高（注：如果您在最初的代碼做一個a = a.tolist()，你得到一個顯著加速）。此外，我使用'mergesort'這是一個穩定的排序，但如果你不需要排序的最終結果，你可以得到一個更快的排序算法。

Answer 2

下面是做這件事的一種方法：

mask = (a > 0) & (a < N) 
elements = a[mask] 
indicies = np.arange(a.size)[mask] 

b = [indicies[elements == i] for i in range(1, N)] 

如果我們兩個時間：

import numpy as np 

a = np.random.randint(0,100,1000) 
N = 10 

def original(a, N): 
    b = [[] for _ in range(N-1)] 
    for indx,elem in enumerate(a): 
     if 0<elem<N: 
      b[elem-1].append(indx) 
    return b 

def new(a, N): 
    mask = (a > 0) & (a < N) 
    elements = a[mask] 
    indicies = np.arange(a.size)[mask] 

    return [indicies[elements == i] for i in range(1, N)] 

「新」 的方式是相當（〜20倍）速度快：

In [5]: %timeit original(a, N) 
100 loops, best of 3: 1.21 ms per loop 

In [6]: %timeit new(a, N) 
10000 loops, best of 3: 57 us per loop 

結果相同：

In [7]: new_results = new(a, N) 

In [8]: old_results = original(a, N) 

In [9]: for x, y in zip(new_results, old_results): 
    ....:  assert np.allclose(x, y) 
    ....: 

In [10]:   

「新」矢量化版本也對更長的序列進行縮放。如果我們對a使用了一個百萬個項目的序列，原始解決方案需要稍微超過1秒，而新版本只需要17毫秒（加速~70倍）。