有效地计算numpy数组中的零元素？

Posted 2023-02-23

【中文标题】有效地计算numpy数组中的零元素？

【英文标题】：Efficiently count zero elements in numpy array?

【发布时间】：2017-08-12 11:34:51

【问题描述】：

我需要计算numpy 数组中零元素的数量。我知道 numpy.count_nonzero 函数，但似乎没有用于计算零元素的模拟。

我的数组不是很大（通常少于 1E5 个元素），但操作执行了几百万次。

我当然可以使用len(arr) - np.count_nonzero(arr)，但我想知道是否有更有效的方法。

这是我目前如何做的 MWE：

import numpy as np
import timeit

arrs = []
for _ in range(1000):
    arrs.append(np.random.randint(-5, 5, 10000))


def func1():
    for arr in arrs:
        zero_els = len(arr) - np.count_nonzero(arr)


print(timeit.timeit(func1, number=10))

【问题讨论】：

count_nonzero 是一个非常基本的编译操作。无论您想知道零的数量还是非零的数量，您仍然需要遍历整个数组。让 numpy 在编译后的代码中做到这一点，不用担心效率。

为什么你认为len(arr) - np.count_nonzero(arr)效率低？

count_nonzero 和潜在的 count_zero 之间的差异是一个减法。这与您将获得的效率差不多。

@juanpa.arrivillaga len(arr) 是通过函数调用进行的属性查找。纯属性查找a.size 花费的时间减少了 25%。

@DYZ 是的，无论如何你都应该使用a.size，特别是因为len(a) 会给多维数组提供错误的答案。但我认为这不是 OP 所指的......

【参考方案1】：

2x 更快的方法是仅使用 np.count_nonzero()，但根据需要使用 条件。

In [3]: arr
Out[3]: 
array([[1, 2, 0, 3],
      [3, 9, 0, 4]])

In [4]: np.count_nonzero(arr==0)
Out[4]: 2

In [5]:def func_cnt():
            for arr in arrs:
                zero_els = np.count_nonzero(arr==0)
                # here, it counts the frequency of zeroes actually

您也可以使用np.where()，但它比np.count_nonzero() 慢

In [6]: np.where( arr == 0)
Out[6]: (array([0, 1]), array([2, 2]))

In [7]: len(np.where( arr == 0))
Out[7]: 2

---

效率：（按降序排列）

In [8]: %timeit func_cnt()
10 loops, best of 3: 29.2 ms per loop

In [9]: %timeit func1()
10 loops, best of 3: 46.5 ms per loop

In [10]: %timeit func_where()
10 loops, best of 3: 61.2 ms per loop

---

使用加速器实现更多加速

如果您可以使用加速器 (GPU/TPU)，现在可以在 JAX 的帮助下实现超过 3 个数量级的速度提升 .使用 JAX 的另一个优点是 NumPy 代码只需很少修改即可使其与 JAX 兼容。下面是一个可重现的例子：

In [1]: import jax.numpy as jnp
In [2]: from jax import jit

# set up inputs
In [3]: arrs = []
In [4]: for _ in range(1000):
   ...:     arrs.append(np.random.randint(-5, 5, 10000))

# JIT'd function that performs the counting task
In [5]: @jit
   ...: def func_cnt():
   ...:     for arr in arrs:
   ...:         zero_els = jnp.count_nonzero(arr==0)

---

# efficiency test
In [8]: %timeit func_cnt()
15.6 µs ± 391 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

【讨论】：

np.where 使用 np.count_nonzero（在编译级别）来确定它返回的数组的大小。

把它放在一个函数调用中以便公平比较......在这种情况下，我得到 1ms 而不是 160us 的 OP 方法......

我试过了，它的剃须速度比func1 使用的时间少了约 10%。

新功能比我原来的功能快 2 倍。我认为这是最好的。谢谢！

在我的计算机上使用 numpy 版本：1.12.1 和列表中的 1000 个数组，我得到 func1：1000 个循环，最好的 3：每个循环 376 µs 和 func_cnt：1000 个循环，最好的 3 ：每个循环 1.65 毫秒。所以我现在很困惑。