python - 将非零 numpy 数组元素附加到列表的最快方法

Question

我想将 numpy 数组 arr 中的所有非零元素添加到列表 out_list 中。之前的研究表明，对于 numpy 数组，使用 np.nonzero 是最有效的。 (下面我自己的基准实际上表明可以使用 np.delete 稍微改进它)。

但是，在我的例子中，我希望我的输出是一个列表，因为我正在组合许多我不知道非零元素数量的数组(因此我无法有效地为它们预分配一个 numpy 数组)。因此，我想知道是否可以利用一些协同作用来加快这一进程。虽然我的简单列表理解方法比纯 numpy 方法慢得多，但我将列表理解与 numba 结合起来得到了一些有希望的结果。

这是我目前的发现:

import numpy as np

n = 60_000  # size of array
nz = 0.3  # fraction of zero elements

arr = (np.random.random_sample(n) - nz).clip(min=0)

# method 1
def add_to_list1(arr, out):
    out.extend(list(arr[np.nonzero(arr)]))

# method 2
def add_to_list2(arr, out):
    out.extend(list(np.delete(arr, arr == 0)))

# method 3
def add_to_list3(arr, out):
    out += [x for x in arr if x != 0]

# method 4 (not sure how to get numba to accept an empty list as argument)
@njit
def add_to_list4(arr):
    return [x for x in arr if x != 0]

out_list = []
%timeit add_to_list1(arr, out_list)

out_list = []
%timeit add_to_list2(arr, out_list)

out_list = []
%timeit add_to_list3(arr, out_list)

_ = add_to_list4(arr)  # call once to compile
out_list = []
%timeit out_list.extend(add_to_list4(arr))

产生以下结果:

2.51 ms ± 137 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
2.19 ms ± 133 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
15.6 ms ± 183 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
1.63 ms ± 158 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

毫不奇怪，numba 优于所有其他方法。其中，方法 2(使用 np.delete)是最好的。我是否遗漏了任何明显的替代方案，这些替代方案利用了我之后转换为列表的事实？你能想出什么来进一步加快这个过程吗？

编辑 1:

.tolist() 的性能:

# method 5
def add_to_list5(arr, out):
    out += arr[arr != 0].tolist()

# method 6
def add_to_list6(arr, out):
    out += np.delete(arr, arr == 0).tolist()

# method 7
def add_to_list7(arr, out):
    out += arr[arr.astype(bool)].tolist()

时间与 numba 相同:

1.62 ms ± 118 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.65 ms ± 104 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each
1.78 ms ± 119 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

编辑2:

这里有一些使用疯狂物理学家建议使用 np.concatenate 来构建 numpy 数组的基准测试。

# construct numpy array using np.concatenate
out_list = []
t = time.perf_counter()
for i in range(100):
    out_list.append(arr[arr != 0])
result = np.concatenate(out_list)
print(f"Time elapsed: {time.perf_counter() - t:.4f}s")

# compare with best list-based method
out_list = []
t = time.perf_counter()
for i in range(100):
    out_list += arr[arr != 0].tolist()
print(f"Time elapsed: {time.perf_counter() - t:.4f}s")

连接 numpy 数组确实产生了另一个显着的加速，尽管它不能直接比较，因为输出是 numpy 数组而不是列表。所以这将取决于精确使用什么是最好的。

Time elapsed: 0.0400s
Time elapsed: 0.1430s

长篇小说;

1/使用 arr[arr != 0] 是所有索引选项中最快的

2/使用 .tolist() 而不是 list(.) 将速度提高 1.3 - 1.5 倍

3/结合 1/和 2/的增益，速度与 numba

相当

4/如果使用 numpy 数组而不是 list 是可以接受的，那么使用 np.concatenate 可以通过与最佳替代方案相比约 3.5 倍

Answer 1

如果您确实在寻找 list 输出，我认为选择的方法是:

def f(arr, out_list):
    out_list += arr[arr != 0].tolist()

它似乎击败了迄今为止在 OP 的问题或其他回复(在撰写本文时)中提到的所有其他方法。

但是，如果您正在寻找作为 numpy 数组的结果，则遵循@MadPhysicist 的版本(略微修改为使用 arr[arr != 0]使用 np.nonzero()) 几乎快 6 倍，请参阅本文末尾。

旁注:我会避免使用%timeit out_list.extend(some_list):它在许多循环中不断添加到out_list 时间。示例:

out_list = []
%timeit out_list.extend([1,2,3])

现在:

>>> len(out_list)
243333333  # yikes

时间

在我机器上的 60K 个项目上，我看到:

out_list = []

a = %timeit -o out_list + arr[arr != 0].tolist()
b = %timeit -o out_list + arr[np.nonzero(arr)].tolist()
c = %timeit -o out_list + list(arr[np.nonzero(arr)])

产量:

1.23 ms ± 10.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.53 ms ± 2.53 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
4.29 ms ± 3.02 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

和:

>>> c.average / a.average
3.476

>>> b.average / a.average
1.244

改为 numpy 数组结果

按照@MadPhysicist，您可以通过不将数组转换为列表，而是使用np.concatenate() 来获得额外的提升:

def all_nonzero(arr_iter):
    """return non zero elements of all arrays as a np.array"""
    return np.concatenate([a[a != 0] for a in arr_iter])

def all_nonzero_list(arr_iter):
    """return non zero elements of all arrays as a list"""
    out_list = []
    for a in arr_iter:
        out_list += a[a != 0].tolist()
    return out_list

from itertools import repeat

ta = %timeit -o all_nonzero(repeat(arr, 100))
tl = %timeit -o all_nonzero_list(repeat(arr, 100))

产量:

39.7 ms ± 107 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
227 ms ± 680 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

和

>>> tl.average / ta.average
5.75