优化数组列分位数识别及布尔数组生成的技术咨询

给定m×n的数组A，以及[0,1]范围内的百分比p，需要生成同尺寸的布尔数组B——当A[i,j]大于等于其所在列的p分位数时，B[i,j]设为True。

原始实现代码 #

import numpy as np

m = 200
n = 300

A = np.random.rand(m, n)

p = 0.3

quant_levels = np.zeros(n)
 
for i in range(n):
    quant_levels[i] = np.quantile(A[:,i],p)
    
B = np.array(A >= quant_levels)

更高效的实现方式 #

numpy的quantile函数支持指定计算轴，完全可以摆脱循环，用向量化操作完成，效率更高：

import numpy as np

m = 200
n = 300

A = np.random.rand(m, n)
p = 0.3

# 按列计算p分位数，keepdims保持维度以便广播
quant_levels = np.quantile(A, p, axis=0, keepdims=True)
# 直接通过广播完成逐元素比较
B = A >= quant_levels

关键说明 #

• 原始循环在列数n较大时会拖慢速度，向量化操作能充分利用numpy的底层优化，大幅提升计算效率。
• axis=0参数指定沿列维度计算分位数，keepdims=True让结果保持(1, n)的形状，这样和(m, n)的A比较时，numpy会自动广播到每行，无需手动调整维度。
• 两种方式生成的B数组结果完全一致，但优化后的代码更简洁，性能更优。

内容的提问来源于stack exchange，提问作者Trevor J Richards