更快地填充列表

Question

更快地填充列表

4

我有一小段代码，用于将列表填充为整数。我需要提高其性能，也许可以将整个代码转换为numpy数组，但我不确定如何做。

这是示例代码：

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)

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

- Gabriel

3

请说明您想要做什么。您的代码似乎应该有一种更简单直接的方法来实现您想要达到的目标。 - Carsten

1

仅供娱乐，排序逻辑的一行代码是 b = [np.where(a == i) for i in range(1, N)]。 - Carsten

@Carsten 该死，错过了那个一行代码，不过在 np.where 后面可能要加一个 [0]？ - deinonychusaur

@deinonychusaur 没错，谢谢，我漏掉了那个。 - Carsten

我想如果你真的非常需要速度，你可以尝试使用Cython来完成，但这似乎只会在目前情况下使事情更加复杂。 - IanH

显示剩余7条评论

2个回答

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' 这个稳定排序算法，但如果你不需要最终结果排序，你可以选择更快的排序算法。

- user2379410

1

我进行了一些测试，确实这个版本比原始版本或Joe的版本快得多。只有在一个情况下，这个版本比原始版本差（使用1000作为a和10000作为N），但是我进行的其余测试都明显更快。对不起Joe，但我将不得不切换到这个版本作为被接受的答案。非常感谢@moarningsun提供的代码！ - Gabriel

1

是的@Gabriel，您原始的“算法”是最有效的。只有通过使用复杂的“算法”绕了一个大弯才能利用Numpy的速度并在适当的大小的a和N上实现总体加速。理想情况下，您将使用原始算法并使用例如Cython或Numba加速它（不确定结果如何）。 - user2379410

1

@Gabriel - 不用担心，这是一个更好的答案！ - Joe Kington

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Joe Kington · Accepted Answer

这里有一种方法可以实现这个功能：

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