二维数组快速索引优化：相似矩阵高频采样性能瓶颈求解

核心优化思路 #

• 抛弃Python原生random模块的循环生成逻辑，改用numpy向量化随机索引生成，利用numpy的C底层实现大幅提速
• 全程基于numpy数组操作，避免将数组转成Python列表的额外开销（原第二种方案的列表转换是性能瓶颈）

具体实现代码 #

import numpy as np

# 假设 similarity_matrix 是 shape=(20000,20) 的 numpy 数组
stream = np.asarray(stream)  # 确保输入流为numpy数组（若原数据是列表则转换）
num_samples = len(stream)

# 直接用numpy生成批量随机列索引，比Python循环快数十倍
random_cols = np.random.randint(0, similarity_matrix.shape[1], size=num_samples)

# 向量化索引一步完成采样，无Python层面循环
sampled_entities = similarity_matrix[stream, random_cols]

额外优化建议 #

• 若输入流是持续批量数据，尽量保持其numpy数组格式，避免反复在列表与数组间转换
• 若采样逻辑允许，可预先生成一批随机索引缓存复用，减少每次生成随机数的开销（需兼顾随机性需求）
• 确保similarity_matrix为连续内存数组，执行similarity_matrix = np.ascontiguousarray(similarity_matrix)优化内存访问效率

性能提升说明 #

• 替换random.randint循环为np.random.randint，随机索引生成速度可提升10~100倍
• 全程numpy向量化操作消除了Python循环开销，整体采样速度至少提升一个数量级

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