如何结合StaticArray与StatsBase降低内存分配？

问题核心在于StatsBase的Weights是可变结构体（mutable struct），每次实例化都会在堆上分配内存；而StaticArray是不可变类型，栈分配无额外开销。要实现无分配地将小型StaticArray转为权重，有两种可行方案：

方案一：自定义静态权重类型（最优） #

直接实现一个基于StaticArray的不可变权重类型，复用StaticArray的栈分配特性，同时兼容StatsBase的接口：

using StaticArrays, StatsBase

# 不可变静态权重类型，持有StaticArray和预计算的权重和
struct StaticWeights{S,T,W<:StaticArray{S,T}} <: AbstractWeights{T}
    weights::W
    sum::T
end

# 自动计算权重和的构造函数
StaticWeights(w::StaticArray) = StaticWeights(w, sum(w))

# 实现StatsBase所需的核心接口
StatsBase.weights(sw::StaticWeights) = sw.weights
Base.sum(sw::StaticWeights) = sw.sum
Base.eltype(::StaticWeights{S,T}) where {S,T} = T
Base.length(::StaticWeights{S}) where {S} = prod(S)
Base.getindex(sw::StaticWeights, i::Int) = sw.weights[i]

使用示例 #

替换原代码中的Weights为StaticWeights即可，构造时无堆分配：

function update_weights_3_static(z, n)
    total = 0.0
    for _ in 1:n
        for zi in z
            w = StaticWeights(zi)
            # 此处可直接用w调用StatsBase函数（如sample、weighted_mean等）
            total += sum(w)
        end
    end
    total
end

# 测试用例：10000个3元素StaticArray
z = [@SVector rand(3) for _ in 1:10000]
@btime update_weights_3_static($z, 1000)  # 分配量接近0

方案二：复用预分配的Weights对象 #

如果必须使用StatsBase原生的Weights类型，可在循环外预分配一个Weights实例，每次仅更新内部的权重数组和总和，避免重复创建对象：

function update_weights_3_reuse(z, n)
    # 假设每个元素是3元素StaticArray，预分配对应大小的数组
    w_arr = Vector{Float64}(undef, 3)
    w = Weights(w_arr)  # 仅创建一次Weights对象
    total = 0.0
    
    for _ in 1:n
        for zi in z
            # 无分配拷贝StaticArray元素到预分配数组
            @inbounds for i in eachindex(zi)
                w_arr[i] = zi[i]
            end
            # 更新Weights的缓存总和（Weights是mutable，可直接修改sum字段）
            w.sum = sum(w_arr)
            # 后续使用w进行计算
            total += sum(w)
        end
    end
    total
end

@btime update_weights_3_reuse($z, 1000)  # 仅在初始化时分配一次

关键原理 #

• 不可变结构体（如StaticWeights）可在栈上分配，无需堆内存开销；而StatsBase的Weights是可变结构体，每次实例化都会产生堆分配。
• 方案一直接复用StaticArray的引用，无需元素拷贝；方案二通过预分配数组避免重复创建Weights对象，仅需拷贝元素（无额外分配）。

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