深圳 网站建设东莞营销网站建设优化

import numpy as np
import timeit

测试数据

开头有 30000 个 0 元素，中间有 100000 个非 0 元素，结尾有 70000 个为 0 元素的 numpy 数组。下面以计算数组末尾的零的数量为例。

array = np.array([0] * 30000 + [i for i in range(1, 1000001)] + [0] * 70000)

方法 1

使用 Python 的循环，从后向前遍历 numpy 数组并统计为 0 元素的个数，当遇到非 0 元素时，则跳出循环。

res1 = 0
for i in range(array.shape[0] - 1, -1, -1):if array[i] == 0:res1 += 1else:break

方法 2

先使用 numpy 的切片器倒序数组，然后使用 numpy 的 consum 的函数计算前缀和；此时除原数组末尾的 0 元素外的其他元素均已变为非 0 元素。于是，先使用 numpy 的 count_nonzero 方法统计满足元素值等于 0 的元素数量。

numpy 的 参考文档：https://numpy.org/doc/stable/reference/generated/numpy.cumsum.html

res2 = np.count_nonzero(np.cumsum(array[::-1]) == 0)

方法 3

在方法 2 的基础上，先直接统计非 0 元素数量，再使用元素总数减去非 0 元素数量得到 0 元素数量。

res3 = array.shape[0] - np.count_nonzero(np.cumsum(array[::-1]))

结果检查

assert res1 == res2 == res3

性能检查

方法 1 计算时间：1.6362712000000001

print(timeit.timeit("res = 0\n""for i in range(array.shape[0] - 1, -1, -1):\n""    if array[i] == 0:\n""        res += 1\n""    else:\n""        break",setup="import numpy as np\n""array = np.array([0] * 30000 + [i for i in range(1, 1000001)] + [0] * 70000)",number=100
))

方法 2 计算时间：0.4067506000000001

因为使用了 numpy 函数替代了 Python 的原生循环，因此性能提升较多。

print(timeit.timeit("np.count_nonzero(np.cumsum(array[::-1]) == 0)",setup="import numpy as np\n""array = np.array([0] * 30000 + [i for i in range(1, 1000001)] + [0] * 70000)",number=100
))

方法 3 计算时间：0.33288949999999984

因为节省了一次计算数组中元素是否等于 0 的判断操作，因此性能有所提升。

print(timeit.timeit("array.shape[0] - np.count_nonzero(np.cumsum(array[::-1]))",setup="import numpy as np\n""array = np.array([0] * 30000 + [i for i in range(1, 1000001)] + [0] * 70000)",number=100
))