基于NumPy实现依赖前序迭代的高效遍历：二进制加1模拟

我想用NumPy模拟任意精度二进制算术，先从简单的加1操作入手。我已经有原生Python的实现代码，能对低位在前、高位在后的二进制数列表执行加1操作（即列表从左到右是二进制数的最低位到最高位）：

def increment_bits(bits):
    """
    返回bits + 1对应的二进制表示，bits是0和1组成的列表
    
    >>> increment_bits([1, 1, 1, 0, 1])   # 23 + 1 == 24
    [0, 0, 0, 1, 1]
    >>> increment_bits([1, 1, 1])         #  7 + 1 ==  8   <- 需要新增一位
    [0, 0, 0, 1]
    """
    new_bits = bits[:]
    for i, v in enumerate(new_bits):
        if v:
            new_bits[i] = 0
        else:
            new_bits[i] = 1
            return new_bits
    # 如果走到这里，说明所有位都是1，需要进位
    new_bits.append(1)
    return new_bits

我的目标是用NumPy实现相同功能并提升速度，但刚接触NumPy，不确定最快速（或最符合NumPy风格）的实现方式。

NumPy实现方案 #

NumPy的核心优势是向量化操作，能避免低效的显式循环，下面是符合NumPy风格的实现：

import numpy as np

def increment_bits_np(bits):
    """
    用NumPy实现二进制数加1，输入为低位在前的0/1数组
    """
    # 复制输入数组，防止修改原数据
    new_bits = bits.copy()
    # 找到所有值为0的位置（从左到右对应二进制数的低位到高位）
    zero_positions = np.where(new_bits == 0)[0]
    
    if zero_positions.size > 0:
        first_zero = zero_positions[0]
        # 将第一个0之前的所有1置为0，把这个0置为1
        new_bits[:first_zero] = 0
        new_bits[first_zero] = 1
        return new_bits
    else:
        # 所有位都是1，生成新数组：原数组全置0后追加高位1
        return np.append(np.zeros_like(new_bits), 1)

关键逻辑说明 #

• 向量化查找：用np.where一次性定位所有0的位置，直接取第一个索引，避免了Python循环逐个判断的开销，长二进制数场景下效率提升明显。
• 批量赋值：new_bits[:first_zero] = 0是批量修改数组元素，比原生Python循环逐个赋值快得多。
• 边界场景处理：当输入全为1时，直接生成新数组，逻辑和原生Python代码的append(1)完全一致。

测试验证 #

# 测试用例1：对应23+1=24
print(increment_bits_np(np.array([1, 1, 1, 0, 1])))  # 输出: [0 0 0 1 1]

# 测试用例2：对应7+1=8
print(increment_bits_np(np.array([1, 1, 1])))  # 输出: [0 0 0 1]

内容的提问来源于stack exchange，提问作者blandish