Julia 中的高效分层随机抽样

Question

我正在尝试编写一个小函数来进行分层随机抽样。也就是说，我对每个元素都有一个组成员资格向量，并且我想为每个组选择一个元素（索引）。因此，输入是所需元素的数量，以及每个元素的组成员资格。输出是一个索引列表。

这是我拥有的功能：

function stratified_sample(n::Int64, groups::Array{Int64})

    # the output vector of indices
    ind = zeros(Int64, n)

    # first select n groups from the total set of possible groups
    group_samp = sample(unique(groups), n, replace = false)

    # cycle through the selected groups
    for i in 1:n
        # for each group, select one index whose group matches the current target group
        ind[i] = sample([1:length(groups)...][groups.==group_samp[i]], 1, replace = false)[1]
    end

    # return the indices
    return ind
end

当我在相对较大的向量上运行此代码时，例如，1000 个不同的组和 40000 个总条目，我得到

julia> groups = sample(1:1000, 40000, replace = true)
40000-element Array{Int64,1}:
 221
 431
 222
 421
 714
 108
 751
 259
   ⋮
 199
 558
 317
 848
 271
 358

julia> @time stratified_sample(5, groups)
  0.022951 seconds (595.06 k allocations: 19.888 MiB)
5-element Array{Int64,1}:
 11590
 17057
 17529
 25103
 20651

并将其与 40000 个可能的五个元素的正常随机抽样进行比较：

julia> @time sample(1:40000, 5, replace = false)
  0.000005 seconds (5 allocations: 608 bytes)
5-element Array{Int64,1}:
 38959
  5850
  3283
 19779
 30063

所以我的代码运行速度慢了近 50k 倍，并且使用了 33k 倍的内存！我到底做错了什么，有没有办法加快这段代码的速度？我的猜测是子集步骤中发生了真正的减速，即[1:length(groups)...][groups.==group_samp[i]]，但我找不到更好的解决方案。

我已经在标准 Julia 包中无休止地搜索此功能，但没有运气。

有什么建议吗？

---

编辑：我已经能够通过随机抽样并检查它是否满足选择 n 个唯一组的要求来加快速度：

function stratified_sample_random(n::Int64, groups::Array{Int64}, group_probs::Array{Float32})
    ind = zeros(Int64, n)
    my_samp = []
    while true
        my_samp = wsample(1:length(groups), group_probs, n, replace = false)
        if length(unique(groups[my_samp])) == n
            break
        end
    end

    return my_samp

end

这里，group_probs 只是一个采样概率向量，其中每个组的元素的总概率为 1/s，其中 s 是该组中元素的数量。例如，如果groups = [1,1,1,1,2,3,3] 对应的概率是group_probs = [0.25, 0.25, 0.25, 0.25, 1, 0.5, 0.5]。这有助于通过最小化选择一组中的多个项目的概率来加快采样速度。总体来说效果还不错：

@time stratified_sample_random(5, groups, group_probs)
  0.000122 seconds (14 allocations: 1.328 KiB)
5-element Array{Int64,1}:
 32209
 10184
 30892
  4861
 30300

通过一些实验，概率加权采样不一定比标准 sample() 快，但这取决于有多少独特的组以及所需的 n 值是多少。

当然，不能保证此函数会随机采样一组唯一的对象，并且它可以永远循环。我的想法是在 while 循环中添加一个计数器，如果它尝试了 10000 次没有运气，那么它将调用我提供的原始 stratified_sample 函数，以确保它返回一个唯一的结果。我不喜欢这种解决方案，而且肯定有更优雅、更简约的方法，但这绝对是一种改进。

Answer 1

在这里，[1:length(groups)...]，您正在喷溅和分配 40000 元素数组 n 次，您应该避免这种情况。这是一个使用范围 inds 的 33 倍快版本。不过，了解实际应用后，我们仍然可以想出一种更快的方法。

function stratified_sample(n::Int64, groups::Array{Int64})

    # the output vector of indices
    ind = zeros(Int64, n)

    # first select n groups from the total set of possible groups
    group_samp = sample(unique(groups), n, replace = false)

    inds = 1:length(groups)
    # cycle through the selected groups
    for i in 1:n
        # for each group, select one index whose group matches the current target group
        ind[i] = sample(inds[groups.==group_samp[i]], 1, replace = false)[1]
    end

    # return the indices
    return ind
end