如何高效获取整数形式随机位串中置1位的索引（类np.nonzero）

方案1：Python内置位运算快速提取置1位 #

通过不断定位并清除最低有效置1位（LSB），只遍历实际置1的位，比全量遍历所有位的循环高效得多：

import numpy as np

def get_set_bit_indices(x, n):
    indices = []
    temp = x
    while temp:
        # 提取最低位的1对应的数值
        lsb = temp & -temp
        # 计算该位的索引（从0开始）
        idx = (lsb).bit_length() - 1
        indices.append(idx)
        # 清除当前最低置1位
        temp ^= lsb
    return indices

# 使用示例
n = 10**6
# 生成n位随机整数位串
bit_int = np.random.bit_generator.randbits(n)
# 获取置1位索引
indices = get_set_bit_indices(bit_int, n)

此方法在1的密度较低时优势尤为明显，无需遍历全部n位。

方案2：NumPy位操作批量处理 #

将整数转换为NumPy位数组后，复用np.nonzero的向量化速度优势，适合需要后续进行NumPy原生操作的场景：

import numpy as np

def int_to_set_bits(x, n):
    # 计算所需字节数，处理n非8倍数的情况
    byte_count = (n + 7) // 8
    # 转换为小端字节数组（保证位索引从0开始对应最低位）
    byte_arr = x.to_bytes(byte_count, byteorder='little')
    # 转为NumPy uint8数组并解包为单个位
    np_bytes = np.frombuffer(byte_arr, dtype=np.uint8)
    bits = np.unpackbits(np_bytes, bitorder='little')
    # 截取前n位，去除补位
    bits = bits[:n]
    # 获取置1位索引
    return np.nonzero(bits)[0]

# 使用示例
n = 10**6
bit_int = np.random.bit_generator.randbits(n)
indices = int_to_set_bits(bit_int, n)

该方法的速度与np.randint + np.nonzero组合相当，同时保留了randbits生成位串的高效性。

方案3：预计算掩码的向量化查询 #

若需多次处理相同长度n的位串，预计算所有位的掩码后批量查询，适合n≤64的场景：

import numpy as np

# 预计算掩码数组（仅需初始化一次）
n = 64
masks = np.left_shift(1, np.arange(n), dtype=np.uint64)

def get_set_bits_with_mask(x, masks):
    is_set = (x & masks) != 0
    return np.nonzero(is_set)[0]

# 使用示例
bit_int = np.random.bit_generator.randbits(n)
indices = get_set_bits_with_mask(bit_int, masks)

注意：当n>64时，需改用分块处理或object dtype，但会损失部分性能。

性能选型建议 #

• 1的密度低且n较大：优先选方案1
• 需要NumPy生态集成：选方案2
• n≤64且需极致速度：选方案3

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