fast rand(::Dict) and rand(::Set) (#18155)

fixes #16231
JuliaLang · Aug 20, 2016 · 4f88b52 · 4f88b52
1 parent 6628674
commit 4f88b52
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/base/random.jl b/base/random.jl
@@ -292,6 +292,19 @@ function rand(r::AbstractRNG, ::Type{Char})
     (c < 0xd800) ? Char(c) : Char(c+0x800)
 end
 
+# random values from Dict or Set (for efficiency)
+function rand(r::AbstractRNG, t::Dict)
+    isempty(t) && throw(ArgumentError("dict must be non-empty"))
+    n = length(t.slots)
+    while true
+        i = rand(r, 1:n)
+        Base.isslotfilled(t, i) && return (t.keys[i] => t.vals[i])
+    end
+end
+rand(t::Dict) = rand(GLOBAL_RNG, t)
+rand(r::AbstractRNG, s::Set) = rand(r, s.dict).first
+rand(s::Set) = rand(GLOBAL_RNG, s)
+
 ## Arrays of random numbers
 
 rand(r::AbstractRNG, dims::Dims) = rand(r, Float64, dims)

diff --git a/test/random.jl b/test/random.jl
@@ -53,6 +53,12 @@ let mt = MersenneTwister()
     rand(coll, 2, 3)
 end
 
+# rand using Dict, Set
+adict = Dict(1=>2, 3=>4, 5=>6)
+@test rand(adict) in adict
+aset = Set(1:10)
+@test rand(aset) in aset
+
 # randn
 @test randn(MersenneTwister(42)) == -0.5560268761463861
 A = zeros(2, 2)