Numpy中沿指定轴寻找最频实数（非整数）的高效实现方法

我们有一个400万行、4列的标准正态分布随机数组：

import numpy as np
x = np.random.randn(4000000, 4)

需求是对每一列（沿axis=0）找出其中最频繁出现的元素。原实现通过循环调用np.histogram逐个维度处理，效率很低——不仅没有利用NumPy的向量化优势，np.histogram本身也有不少额外开销。下面是几种高效优化方案：

原低效实现回顾 #

# 循环构建每个维度的直方图
hist = [np.histogram(x[:, i], bins=200) for i in range(x.shape[-1])]
# 拆分频次和分箱边界
freq, val = (np.array([o[0] for o in hist]).T, np.array([o[1] for o in hist]).T)
# 找最大频次的分箱索引
most_freq = np.argmax(freq, axis=0)
# 提取对应分箱的近似值
result = np.diagonal(val[most_freq])

结果近似为[0., 0., 0., 0.]（符合正态分布的峰值特性），但速度远未达标。

优化方案1：全向量化分箱与统计 #

通过一次性对所有列做离散化和频次统计，替代循环调用np.histogram：

# 1. 获取所有列的极值，生成等宽分箱
col_mins = x.min(axis=0)
col_maxs = x.max(axis=0)
bins = np.linspace(col_mins, col_maxs, 201).T  # shape: (4, 201)

# 2. 向量化离散化所有列，得到每个元素的分箱索引
bin_indices = np.digitize(x, bins, right=True)
# 修正边界值（避免超出分箱范围）
bin_indices[bin_indices == 200] = 199

# 3. 对每列统计分箱频次
freqs = np.array([np.bincount(col, minlength=200) for col in bin_indices.T])

# 4. 找到最大频次的分箱，计算对应中点值
most_freq_idx = np.argmax(freqs, axis=1)
bin_midpoints = (bins[:, :-1] + bins[:, 1:]) / 2
result = bin_midpoints[np.arange(4), most_freq_idx]

np.digitize支持对每列使用独立分箱，一次性完成所有列的离散化；np.bincount比np.histogram更轻量，统计效率更高。

优化方案2：Numba并行加速循环 #

如果需要更极致的性能，用Numba对循环做JIT编译，利用多核并行处理：

import numba

@numba.jit(nopython=True, parallel=True)
def compute_most_freq(x, num_bins=200):
    n_cols = x.shape[1]
    result = np.zeros(n_cols)
    # 并行遍历每一列
    for i in numba.prange(n_cols):
        col = x[:, i]
        col_min = col.min()
        col_max = col.max()
        bin_width = (col_max - col_min) / num_bins
        
        # 计算分箱索引并修正边界
        bin_indices = ((col - col_min) / bin_width).astype(np.int64)
        bin_indices[bin_indices == num_bins] = num_bins - 1
        
        # 统计频次并找峰值分箱
        freq = np.bincount(bin_indices, minlength=num_bins)
        max_idx = np.argmax(freq)
        
        # 计算峰值对应的近似值
        result[i] = col_min + (max_idx + 0.5) * bin_width
    return result

# 调用函数
result = compute_most_freq(x)

Numba的parallel=True会自动将循环分配到多个CPU核心，相比纯Python循环速度提升10倍以上，同时避免了np.histogram的不必要开销。

优化方案3：利用数据特性预先设定分箱 #

因为数据是标准正态分布，99.9%的样本集中在[-4, 4]区间，我们可以直接预先设定分箱范围，省去计算每列极值的时间：

# 预先针对正态分布设定分箱
bins = np.linspace(-4, 4, 201)
# 离散化所有列
bin_indices = np.digitize(x, bins, right=True)
# 修正超出范围的少数样本
bin_indices[bin_indices < 0] = 0
bin_indices[bin_indices >= 200] = 199

# 统计频次并计算结果
freqs = np.array([np.bincount(col, minlength=200) for col in bin_indices.T])
most_freq_idx = np.argmax(freqs, axis=1)
bin_midpoints = (bins[:-1] + bins[1:]) / 2
result = bin_midpoints[most_freq_idx]

这种方案省去了全局极值计算的步骤，对于已知分布的数据效率最优。

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