如何在Julia中高效计算向量分组后的行和与列和？

核心思路 #

直接基于原向量的索引逻辑遍历累加，完全跳过reshape创建中间矩阵的步骤，仅分配结果数组的内存，同时通过一次遍历完成列和与行和的计算，最大化速度与内存效率。

优化实现（以Julia为例） #

function compute_group_sums(vec, group_size)
    n = length(vec)
    # 确保向量长度能被分组大小整除
    @assert n % group_size == 0 "Vector length must be divisible by group size"
    
    n_groups = n ÷ group_size
    # 初始化结果数组，仅分配必要内存
    res_col = similar(vec, n_groups)
    res_row = similar(vec, group_size)
    fill!(res_col, 0)
    fill!(res_row, 0)
    
    # 一次遍历完成两个求和计算
    for i in 1:n
        group_idx = (i - 1) ÷ group_size + 1
        row_idx = (i - 1) % group_size + 1
        res_col[group_idx] += vec[i]
        res_row[row_idx] += vec[i]
    end
    
    return res_col, res_row
end

用法与验证 #

test_vector = 1:100
group_size = 10

# 优化方法计算
res_col, res_row = compute_group_sums(test_vector, group_size)

# 原方法用于对比验证
reshape_mat = reshape(test_vector, group_size, :)
res_1 = sum(reshape_mat, dims=1) |> vec
res_2 = sum(reshape_mat, dims=2) |> vec

# 确认结果一致
@assert res_col == res_1
@assert res_row == res_2

性能优势说明 #

1. 内存控制：仅分配两个结果数组的内存，完全避免reshape带来的中间矩阵（哪怕是视图类型的中间结构也会引入额外的抽象开销）。
2. 速度提升：一次遍历完成两类求和，减少了数据遍历次数；直接的循环逻辑在Julia中会被编译为高效的机器码，比reshape+sum的组合少了多层抽象的性能损耗。
3. 可复用性：封装为函数后，多次调用时不会重复产生中间内存分配，适合高频调用场景。

额外优化提示 #

• 如果向量元素类型固定（如Int/Float64），可以直接用zeros(Int, n_groups)替代similar，进一步简化初始化逻辑。
• 针对超大规模向量，可结合多线程并行遍历（需确保线程安全，比如用Threads.@threads配合原子操作或分段累加）。
• 若分组大小是固定常量，可将函数改写为生成函数（如compute_group_sums(::Type{Val{10}}, vec)），让编译器做更多静态优化。

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