如何优化Python中生成二进制向量的for循环执行速度

原有代码的核心性能瓶颈：

• 直接对序列做in查询时间复杂度为O(n)，叠加遍历词汇表的循环后整体复杂度为O(词汇表长度 × 输入序列长度)，词汇量大时耗时会陡增
• 存在numpy数组转列表、再转回numpy数组的冗余操作，增加了额外开销
• Python原生for循环执行效率远低于numpy向量化操作

最小改动优化版本 #

改动最小，无需调整类其他逻辑，即可获得数倍性能提升：

def objects_to_binary_vector(self, object_seq: Sequence[Any]) -> np.ndarray:
    # 转集合将查询复杂度降到O(1)
    obj_set = set(object_seq)
    # 直接初始化numpy数组，省去列表转换开销
    binary = np.zeros(len(self.vocab) + self.start, dtype=np.int8)
    for idx, word in enumerate(self.vocab):
        if word in obj_set:
            binary[idx + self.start] = 1
    return binary

高性能向量化版本 #

适合词汇规模较大的场景，可获得数十倍甚至上百倍的性能提升：

def objects_to_binary_vector(self, object_seq: Sequence[Any]) -> np.ndarray:
    binary = np.zeros(len(self.vocab) + self.start, dtype=np.int8)
    # 【建议提前在类初始化阶段生成该映射，不要每次调用都重复计算】
    if not hasattr(self, 'vocab_to_idx'):
        self.vocab_to_idx = {word: idx + self.start for idx, word in enumerate(self.vocab)}
    # 批量提取存在的元素索引，一次性赋值
    exist_ids = [self.vocab_to_idx[obj] for obj in object_seq if obj in self.vocab_to_idx]
    binary[exist_ids] = 1
    return binary

如果词汇表是固定不变的，把vocab_to_idx的生成逻辑挪到__init__方法里，性能会进一步提升。

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