Efficient Stratified Random Sampling in Julia

I am trying to write a little function to do stratified random sampling. That is, I have a vector of group membership for each element, and I want to select a single element (index) for each group. Thus, the input is the number of elements desired, and the group membership for each element. The output is a list of indices.

Here's the function I have:

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

When I run this code on a relatively large vector, for example, 1000 different groups and 40000 total entries, I get

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

And to compare this to normal random sampling of five elements from 40000 possible:

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

So my code is running nearly 50k times slower and using up 33k times more memory! What in world have I done wrong, and is there a way to speed up this code? My guess is the real slow down is happening in the subsetting step, i.e., [1:length(groups)...][groups.==group_samp[i]], but I can't find a better solution.

I've searched endlessly for this function within standard Julia packages, but no luck.

Any suggestions?

EDIT: I've been able to speed it up a lot by just taking a random sample, and checking to see if it satisfies the requirement that there are n unique groups selected:

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

Here, group_probs is just a vector of sampling probabilities, where each group's elements have a total probability of 1/s, where s is the number of elements in that group. For example, if groups = [1,1,1,1,2,3,3] the the corresponding probabilities would be group_probs = [0.25, 0.25, 0.25, 0.25, 1, 0.5, 0.5]. This helps speed up the sampling by minimizing the probability of selecting multiple items of one group. Overall it works fairly well:

@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

From a bit of experimenting, weighted sampling by probability is not necessarily faster than standard sample(), but this depends on how many unique groups and what the desired n value is.

Of course, there's no guarantee this function will randomly sample a unique set of objects, and it could cycle forever. My thought is to add a counter to the while loop, and if it's tried for something like 10000 times with no luck, then it will call the original stratified_sample function I presented, to ensure it returns a unique result. I don't love this solution, and there must be a more elegant and parsimonious approach, but it's definitely an improvement.