From fbf6a8e785b2dec9b67d78f7a34c748091ff098f Mon Sep 17 00:00:00 2001
From: "Steven G. Johnson" <stevenj@alum.mit.edu>
Date: Fri, 2 May 2014 22:25:43 -0400
Subject: [PATCH] faster sprand, new randsubseq (fixes #6714 and #6708)

---
 base/abstractarray.jl       | 40 +++++++++++++++++++++++
 base/exports.jl             |  2 ++
 base/sparse/sparsematrix.jl | 65 +++++++++++++------------------------
 doc/stdlib/base.rst         | 13 ++++++++
 4 files changed, 78 insertions(+), 42 deletions(-)

diff --git a/base/abstractarray.jl b/base/abstractarray.jl
index b11e1f5d0a0431..161c29ce1df7e8 100644
--- a/base/abstractarray.jl
+++ b/base/abstractarray.jl
@@ -1314,3 +1314,43 @@ push!(A, a, b, c...) = push!(push!(A, a, b), c...)
 unshift!(A) = A
 unshift!(A, a, b) = unshift!(unshift!(A, b), a)
 unshift!(A, a, b, c...) = unshift!(unshift!(A, c...), a, b)
+
+# Fill S (resized as needed) with a random subsequence of A, where
+# each element of A is included in S with independent probability p.
+# (Note that this is different from the problem of finding a random
+#  size-m subset of A where m is fixed!)
+function randsubseq!(S::AbstractArray, A::AbstractArray, p::Real)
+    0 <= p <= 1 || throw(ArgumentError("probability $p not in [0,1]"))
+    n = length(A)
+    p == 1 && return copy!(resize!(S, n), A)
+    empty!(S)
+    p == 0 && return S
+    sizehint(S, int(p * length(A)))
+    i = 0
+    if p > 0.15 # empirical threshold for trivial O(n) algorithm to be better
+        for j = 1:n
+            rand() <= p && push!(S, A[j])
+        end
+    else
+        # Skip through A, in order, from each element i to the next element i+s
+        # included in S. The probability that the next included element is 
+        # s==k (k > 0) is (1-p)^(k-1) * p, and hence the probability (CDF) that
+        # s is in {1,...,k} is 1-(1-p)^k = F(k).   Thus, we can draw the skip s
+        # from this probability distribution via the discrete inverse-transform
+        # method: s = iceil(F^{-1}(u)) where u = rand(), which is simply
+        # s = iceil(log(rand()) / log1p(-p)).
+        L = 1 / log1p(-p)
+        while true
+            s = log(rand()) * L # note that rand() < 1, so s > 0
+            s >= n - i && return S # compare before iceil to avoid overflow
+            push!(S, A[i += iceil(s)])
+        end
+        # [This algorithm is similar in spirit to, but much simpler than,
+        #  the one by Vitter for a related problem in "Faster methods for
+        #  random sampling," Comm. ACM Magazine 7, 703-718 (1984).]
+    end
+    return S
+end
+
+randsubseq{T}(A::AbstractArray{T}, p::Real) = randsubseq!(T[], A, p)
+
diff --git a/base/exports.jl b/base/exports.jl
index 55ebf638c4b43d..9d4e47d79ddf22 100644
--- a/base/exports.jl
+++ b/base/exports.jl
@@ -545,6 +545,8 @@ export
     promote_shape,
     randcycle,
     randperm,
+    randsubseq!,
+    randsubseq,
     range,
     reducedim,
     repmat,
diff --git a/base/sparse/sparsematrix.jl b/base/sparse/sparsematrix.jl
index 07c880fe5d8b49..c1f7b584ac49cc 100644
--- a/base/sparse/sparsematrix.jl
+++ b/base/sparse/sparsematrix.jl
@@ -335,56 +335,37 @@ function findnz{Tv,Ti}(S::SparseMatrixCSC{Tv,Ti})
     return (I, J, V)
 end
 
-truebools(n::Integer) = ones(Bool, n)
-function sprand{T}(m::Integer, n::Integer, density::FloatingPoint, rng::Function,::Type{T}=eltype(rng(1)))
-    0 <= density <= 1 || throw(ArgumentError("density must be between 0 and 1"))
+function sprand{T}(m::Integer, n::Integer, density::FloatingPoint,
+                   rng::Function,::Type{T}=eltype(rng(1)))
+    0 <= density <= 1 || throw(ArgumentError("$density not in [0,1]"))
     N = n*m
     N == 0 && return spzeros(T,m,n)
     N == 1 && return rand() <= density ? sparse(rng(1)) : spzeros(T,1,1)
-    # if density < 0.5, we'll randomly generate the indices to set
-    #        otherwise, we'll randomly generate the indices to skip
-    K = (density > 0.5) ? N*(1-density) : N*density
-    # Use Newton's method to invert the birthday problem
-    l = log(1.0-1.0/N)
-    k = K
-    k = k + ((1-K/N)*exp(-k*l) - 1)/l
-    k = k + ((1-K/N)*exp(-k*l) - 1)/l # for K<N/2, 2 iterations suffice
-    ik = int(k)
-    ind = sort(rand(1:N, ik))
-    uind = Array(Int, 0)   # unique indices
-    sizehint(uind, int(N*density))
-    if density < 0.5
-        if ik == 0
-            return sparse(Int[],Int[],Array(T,0),m,n)
-        end
-        j = ind[1]
-        push!(uind, j)
-        uj = j
-        for i = 2:length(ind)
-            j = ind[i]
-            if j != uj
-                push!(uind, j)
-                uj = j
-            end
-        end
-    else
-        push!(ind, N+1) # sentinel
-        ii = 1
-        for i = 1:N
-            if i != ind[ii]
-                push!(uind, i)
-            else
-                while (i == ind[ii])
-                    ii += 1
-                end
-            end
+
+    I, J = Array(Int, 0), Array(Int, 0) # indices of nonzero elements
+    sizehint(I, int(N*density))
+    sizehint(J, int(N*density))
+
+    # density of nonzero columns:
+    coldensity = -expm1(m*log1p(-density)) # = 1 - (1-density)^m 
+    
+    rows = Array(Int, 0)
+    for j in randsubseq(1:n, coldensity)
+        randsubseq!(rows, 1:m, density)
+        append!(I, rows)
+        nrows = length(rows)
+        Jlen = length(J)
+        resize!(J, Jlen+nrows)
+        for i = Jlen+1:length(J)
+            J[i] = j
         end
     end
-    I, J = ind2sub((m,n), uind)
-    return sparse_IJ_sorted!(I, J, rng(length(uind)), m, n, +)  # it will never need to combine
+    return sparse_IJ_sorted!(I, J, rng(length(I)), m, n, +)  # it will never need to combine
 end
+
 sprand(m::Integer, n::Integer, density::FloatingPoint) = sprand(m,n,density,rand,Float64)
 sprandn(m::Integer, n::Integer, density::FloatingPoint) = sprand(m,n,density,randn,Float64)
+truebools(n::Integer) = ones(Bool, n)
 sprandbool(m::Integer, n::Integer, density::FloatingPoint) = sprand(m,n,density,truebools,Bool)
 
 spones{T}(S::SparseMatrixCSC{T}) =
diff --git a/doc/stdlib/base.rst b/doc/stdlib/base.rst
index 7a1aec99e3d813..e0e327fb896c85 100644
--- a/doc/stdlib/base.rst
+++ b/doc/stdlib/base.rst
@@ -3814,6 +3814,19 @@ Indexing, Assignment, and Concatenation
 
    Throw an error if the specified indexes are not in bounds for the given array.
 
+.. function:: randsubseq(A, p) -> Vector
+   
+   Return a vector consisting of a random subsequence of the given array ``A``,
+   where each element of ``A`` is included (in order) with independent
+   probability ``p``.   (Complexity is linear in ``p*length(A)``, so this
+   function is efficient even if ``p`` is small and ``A`` is large.)
+
+.. function:: randsubseq!(S, A, p)
+
+   Like ``randsubseq``, but the results are stored in ``S`` (which is
+   resized as needed).
+   
+
 Array functions
 ~~~~~~~~~~~~~~~