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

朱莉娅的高效分层随机抽样

1 个答案: