(Julia) 快速地对一个数组进行列求和

Question

我有一个大小为MxN的浮点数数组A，我想进行以下操作：
对于A的每一列，统计小于某个阈值（例如0.5）的元素数量。
Julia中最快的方法是使用零初始化结果向量，然后按列遍历数组A并根据需要增加结果向量。使用for循环很容易实现，例如：

function count(A)
    (cols, rows) = size(A)
    result = fill(0, cols)
    for j in 1:rows
        for i in 1:cols
            result[i] += A[i,j] <= 0.5
        end
    end
end

这样遍历A时，它的顺序与内存中的布局相同，并且不会分配不必要的额外空间。但是，我不确定如何使用广播运算符 .<=0.5、sum 等来实现这一点。一种方法是

sum(A .<= 0.5, dims=1)

但这样分配新的内存来执行操作，速度大约比计数函数慢2-3倍（我测试了高达8000x8000的数组）。是否可能在只使用少量代码（例如第二种方法）的情况下实现计数函数的性能？

- Matias Frank Jensen

Answer 1

表达您的条件的一种自然方式是：

count.(<=(0.5), eachcol(A))

或者

vec(sum(<=(0.5), A, dims=1))

后者比前者稍慢，但分配的内存要少得多，所以性能平衡可能取决于矩阵的维度。

总的来说，两者都应该相当快速和易于阅读。

顺便提一下 - 你的代码在行而不是列中执行计数。这是将其更正为在列中执行计数的方法：

function count2(A)
    (rows, cols) = size(A)
    result = fill(0, cols)
    for i in 1:cols
        for j in 1:rows
            result[i] += A[j,i] <= 0.5
        end
    end
    result
end

最后注意，sum(A .<= 0.5, dims=1) 执行的操作与其他函数不同，因为它返回一个 Matrix 而不是一个 Vector。

- Bogumił Kamiński