更快地填写列表答案

【问题标题】：Filling a list faster更快地填写列表
【发布时间】：2014-03-20 10:34:30
【问题描述】：

我有一小段代码用于填充整数列表。我需要改进它的性能，也许将整个东西转换成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)

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

【问题讨论】：

请说明你想做什么。您的代码看起来应该是一种更直接的方式来完成您想要实现的目标。
您可能需要生成器或 itertools。但我同意@Carsten。解释你的目标是什么。顺便说一句 b 可能是 defaultdict(list)
只是为了好玩，排序逻辑的单行是b = [np.where(a == i) for i in range(1, N)]。
@Carsten 我可以解释这个块的用途，但需要几页。这是一小块隐藏在我正在处理的更大（数千行）代码中的代码块，它做了一件非常小的事情。这就是为什么我忽略了所有会使它不必要地复杂化的东西，并尽可能简单地发布它。有关更多上下文，我来自这个问题stackoverflow.com/questions/21856907/…，我现在正在尝试提高 Add 2 中提到的函数的性能。
@Carsten 该死的错过了那条线，但可能会在 np.where 上添加一个 [0] 吗？

标签： python arrays performance list numpy

【解决方案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 倍）。

【讨论】：

很好的答案@Joe，非常感谢。只是检查一下，您能否确认对于较大的 N 值（比如 a = np.random.randint(0,100,5000) 和 N=1000），这个答案实际上比原始代码慢一点？
@Gabriel - 哎呀，是的，你是对的。对于N 的大值和a 的长度，对于低千位的原始函数，我的时间会稍快一些。不过，对于很长的a，“新”功能再次变得更快。不过，我看不到任何进一步加速的方法。祝你好运！

【解决方案2】：

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

【讨论】：

我做了一些测试，确实这个版本比原来的版本或乔的版本快得多。只有在一种情况下，这个版本比原来的版本更差（使用1000 表示a 和10000` 表示N）但我执行的其余测试明显更快。对不起，乔，但我必须切换到这个作为接受的答案。非常感谢@moarningsun 的代码！
是的@Gabriel，你原来的“算法”简直是最有效的。只有通过使用复杂的“算法”，您才能充分利用 Numpy 的速度并在a 和N 的正确大小下实现整体加速。理想情况下，您将使用原始算法并使用例如加速它Cython 或 Numba（但不确定结果会如何）。
@Gabriel - 不用担心，这是一个更好的答案！