当前位置：首页 > news >正文

【Python性能优化实例】计算 numpy 数组首尾为 0 的数量

news 来源：原创 2025/5/1 23:31:39

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 += 1
    else:
        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
))