计算Numba中的numpy数组中非零值的数量

问题描述

非常简单.我正在尝试计算用Numba(njit())编译的NumPy jit中数组中非零值的数量. Numba禁止我尝试以下操作.

Very simple. I am trying to count the number of non-zero values in an array in NumPy jit compiled with Numba (njit()). The following I've tried is not allowed by Numba.

1. a[a != 0].size
2. np.count_nonzero(a)
3. len(a[a != 0])
4. len(a) - len(a[a == 0])

1. a[a != 0].size
2. np.count_nonzero(a)
3. len(a[a != 0])
4. len(a) - len(a[a == 0])

如果仍然有一种更快，更pythonic且优雅的方式，我不想用于循环.

I don't want to use for loops if there is still a faster, more pythonic and elegant way.

对于想要查看完整代码示例的评论者...

For that commenter that wanted to see a full code example...

import numpy as np
from numba import njit

@njit()
def n_nonzero(a):
    return a[a != 0].size

推荐答案

您也可以考虑计算非零值:

You may also consider, well, counting the nonzero values:

import numba as nb

@nb.njit()
def count_loop(a):
    s = 0
    for i in a:
        if i != 0:
            s += 1
    return s

我知道这似乎是错的，但请忍受:

I know it seems wrong, but bear with me:

import numpy as np
import numba as nb

@nb.njit()
def count_loop(a):
    s = 0
    for i in a:
        if i != 0:
            s += 1
    return s

@nb.njit()
def count_len_nonzero(a):
    return len(np.nonzero(a)[0])

@nb.njit()
def count_sum_neq_zero(a):
    return (a != 0).sum()

np.random.seed(100)
a = np.random.randint(0, 3, 1000000000, dtype=np.uint8)
c = np.count_nonzero(a)
assert count_len_nonzero(a) == c
assert count_sum_neq_zero(a) == c
assert count_loop(a) == c

%timeit count_len_nonzero(a)
# 5.94 s ± 141 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit count_sum_neq_zero(a)
# 848 ms ± 80.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit count_loop(a)
# 189 ms ± 4.41 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

实际上它比np.count_nonzero快，由于某些原因，它可能变得很慢:

It is in fact faster than np.count_nonzero, which can get quite slow for some reason:

%timeit np.count_nonzero(a)
# 4.36 s ± 69.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

这篇关于计算Numba中的numpy数组中非零值的数量的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！