CODEWITHSUNDEEP
randomjuliasamplesampling
xiaodai
xiaodai

Reputation: 16004

Julia: best way to sample from successively shrinking range?

I would like to sample k numbers where the first number is sampled from 1:n and the second from 1:n-1 and the third from 1:n-2 and so on.

I have the below implementation

function shrinksample(n,k)
    [rand(1:m) for m in n:-1:n-k+1]
end

Are there faster solutions in Julia?

Upvotes: 1

Views: 144

Answers (2)

AboAmmar
AboAmmar

Reputation: 5559

If this is not very sensitive to randomness (you're not doing cryptography), the following should be amazingly fast and very simple:

blazingshrinksample(n,k) = (Int)[trunc(Int,(n-m)rand()+1) for m in 0:k-1]

Testing this along with your implementation and with Dan's, I got this:

using BenchmarkTools

@btime shrinksample(10000,10000);
  259.414 μs (2 allocations: 78.20 KiB)

@btime fastshrinksample(10000,10000);
  66.713 μs (2 allocations: 78.20 KiB)

@btime blazingshrinksample(10000,10000);
  33.614 μs (2 allocations: 78.20 KiB)

Upvotes: 1

Dan Getz
Dan Getz

Reputation: 18217

The following takes ideas from the implementation of randperm and since n and k are of the same order, this is appropriate as the same type of randomness is needed (both have output space of size n factorial):

function fastshrinksample(r::AbstractRNG,n,k)
    a = Vector{typeof(n)}(k)
    @assert n <= Int64(2)^52
    k == 0 && return a
    mask = (1<<(64-leading_zeros(n)))-1
    nextmask = mask>>1
    nn = n
    for i=1:k
        a[i] = 1+Base.Random.rand_lt(r, nn, mask)
        nn -= 1
        if nn == nextmask
            mask, nextmask = nextmask, nextmask>>1
        end
    end
    return a
end

fastshrinksample(n,k) = fastshrinksample(Base.Random.GLOBAL_RNG, n, k)

Benchmarking gives a 3x improvement for one tested instance:

julia> using BenchmarkTools

julia> @btime shrinksample(10000,10000);
  310.277 μs (2 allocations: 78.20 KiB)

julia> @btime fastshrinksample(10000,10000);
  91.815 μs (2 allocations: 78.20 KiB)

The trick is mainly to use the internal Base.Random.rand_lt instead of regular rand(1:n)

Upvotes: 4

Related Questions