基于Python/Numba快速实现区间一维直方图的方案咨询

需求概述 #

基于两组无序边界定义的区间、对应权重与指定步长，快速生成一维直方图，允许通过Cython、Numba等工具进行性能优化，输出需包含直方图数组、全局最小值与最大值。

示例数据 #

import numpy as np

boundaries1 = np.array([511.4, 517.5, 517.2, 520.1, 519.8])
boundaries2 = np.array([517.3, 517.1, 517.8, 518.2, 515.0])
weights = np.array([0.2, 0.4, 0.6, 0.8, 1])
stepsize = 0.1

现有Numba实现 #

from numba import njit
import numpy as np

@njit(cache=True)
def create_histogram(boundaries1, boundaries2, weights, stepsize):
    minimum = np.minimum(np.min(boundaries1),np.min(boundaries2))
    maximum = np.maximum(np.max(boundaries1),np.max(boundaries2))

    num_bins = int(np.round((maximum-minimum+stepsize)/stepsize))
    hist = np.zeros(num_bins, dtype=np.float64)
        
    lower = np.round((np.minimum(boundaries1, boundaries2)-minimum)/stepsize).astype(np.int32)
    upper = np.round((np.maximum(boundaries1, boundaries2)-minimum)/stepsize).astype(np.int32)
    
    for i in range(lower.shape[0]):
        hist[lower[i]:upper[i]+1] += weights[i]
    
    return hist, minimum, maximum

优化思路：向量化差分替代循环累加 #

原实现中对每个区间循环累加权重的操作，可通过差分法+前缀和实现完全向量化，避免显式循环，结合Numba或纯Numpy都能获得显著性能提升，尤其在区间数量庞大时效果更明显：

核心逻辑 #

1. 计算全局最小/最大值与分箱数，和原逻辑一致
2. 构建差分数组：对每个区间的起始位置lower加上对应权重，对upper+1位置减去对应权重（注意边界处理，避免越界）
3. 对差分数组做前缀和，得到最终直方图

优化实现（Numba加速版） #

from numba import njit
import numpy as np

@njit(cache=True)
def create_histogram_opt(boundaries1, boundaries2, weights, stepsize):
    minimum = np.minimum(np.min(boundaries1), np.min(boundaries2))
    maximum = np.maximum(np.max(boundaries1), np.max(boundaries2))
    
    num_bins = int(np.round((maximum - minimum + stepsize) / stepsize))
    delta = np.zeros(num_bins + 1, dtype=np.float64)  # 多留一个位置处理upper+1的边界
    
    lower = np.round((np.minimum(boundaries1, boundaries2) - minimum) / stepsize).astype(np.int32)
    upper = np.round((np.maximum(boundaries1, boundaries2) - minimum) / stepsize).astype(np.int32)
    
    # 向量化更新差分数组（Numba会自动优化这部分循环）
    for i in range(lower.shape[0]):
        delta[lower[i]] += weights[i]
        if upper[i] + 1 <= num_bins:
            delta[upper[i] + 1] -= weights[i]
    
    # 前缀和生成直方图
    hist = np.cumsum(delta[:-1])
    return hist, minimum, maximum

纯Numpy向量化实现（无需Numba） #

如果不需要兼容旧环境，纯Numpy的np.add.at可以直接完成差分数组的更新，完全避免Python循环：

import numpy as np

def create_histogram_np(boundaries1, boundaries2, weights, stepsize):
    minimum = np.minimum(boundaries1.min(), boundaries2.min())
    maximum = np.maximum(boundaries1.max(), boundaries2.max())
    
    num_bins = int(np.round((maximum - minimum + stepsize) / stepsize))
    delta = np.zeros(num_bins + 1, dtype=np.float64)
    
    lower = np.round((np.minimum(boundaries1, boundaries2) - minimum) / stepsize).astype(np.int32)
    upper = np.round((np.maximum(boundaries1, boundaries2) - minimum) / stepsize).astype(np.int32)
    
    # 批量更新差分数组
    np.add.at(delta, lower, weights)
    valid_upper = upper + 1
    mask = valid_upper <= num_bins
    np.add.at(delta, valid_upper[mask], -weights[mask])
    
    hist = np.cumsum(delta[:-1])
    return hist, minimum, maximum

示例结果 #

运行示例数据后，得到：

• 全局最小值：511.4
• 全局最大值：520.1
• 直方图片段（部分关键位置）：
  • 第0-58个分箱（对应区间511.4~517.2）：权重累计0.2
  • 第59个分箱（517.3）：累计0.2+0.4+1=1.6
  • 第60个分箱（517.4）：累计1.6+0.6=2.2
  • ...后续分箱根据区间叠加权重

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