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

Question

我需要统计 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))

- Gabriel

---

count_nonzero是一个非常基本的编译操作。无论您想知道零的数量还是非零的数量，您仍然必须遍历整个数组。让numpy在编译代码中完成这项工作，不用担心效率问题。 - hpaulj

你为什么认为 len(arr) - np.count_nonzero(arr) 是低效的？ - juanpa.arrivillaga

count_nonzero和潜在的count_zero之间的区别在于一次减法。这是你能得到的最高效的方式。 - pvg

你知道len(are)只是一个简单的属性查找，对吧？它不会再次迭代数组... - juanpa.arrivillaga

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

@DYZ 是的，你应该无论如何使用 a.size，特别是对于多维数组，而 len(a) 会给出错误的答案。但我认为这不是 OP 所指的问题…… - juanpa.arrivillaga

Answer 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

---

加速器可带来更多速度提升

使用JAX并配备加速器（GPU/TPU），现在可以实现超过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)

- kmario23

---

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

新函数的速度比我的原始函数快2倍，我认为这已经是最好的了。谢谢！ - Gabriel

在我的电脑上，numpy版本为1.12.1，列表中有1000个数组。对于func1函数，我得到了1000次循环，3次中最佳结果为每次376微秒；而对于func_cnt函数，我得到了1000次循环，3次中最佳结果为每次1.65毫秒。所以现在我感到困惑。 - plasmon360

以防其他人有我这样的想法：我想尝试使用zero_els = np.sum(arr == 0)，但实际上它比上面描述的任何方法都要慢。 - Brunox13

使用一个名为'count_nonzero'的函数来计算零的出现次数。很好。非常好。我不会在这里做任何改变。这是很好的命名，而且很有意义。 - Jarvis