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)

这就是它的作用：

表示integ

a

查看它是否位于给定范围(0,N)
如果是，则将其索引存储在b的子列表中，其中所述子列表的索引是原始整数减1（integ-1）

这段代码运行得非常快，但我的实际代码使用了更大的列表，因此需要提高其性能。

Answer 1

这是一种方法：

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倍）。

Answer 2

试试这个解决方案！上半场我无耻地从乔的答案中偷走了，但之后它使用了排序和二分搜索，它可以更好地与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'这是一个稳定的排序，但如果你不需要对最终结果进行排序，你就可以使用更快的排序算法。

更快地填写列表

2 个答案: