Allow quantile function to accept tuples

Fixes #18458
JuliaLang · May 4, 2017 · dd8e6dd · dd8e6dd
1 parent 072317b
commit dd8e6dd
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 19 deletions.
diff --git a/base/statistics.jl b/base/statistics.jl
@@ -648,19 +648,10 @@ function quantile!(q::AbstractArray, v::AbstractVector, p::AbstractArray;
     if size(p) != size(q)
         throw(DimensionMismatch("size of p, $(size(p)), must equal size of q, $(size(q))"))
     end
+    isempty(q) && return q
 
-    isempty(v) && throw(ArgumentError("empty data vector"))
-
-    lv = length(v)
-    if !sorted
-        minp, maxp = extrema(p)
-        lo = floor(Int,1+minp*(lv-1))
-        hi = ceil(Int,1+maxp*(lv-1))
-
-        # only need to perform partial sort
-        sort!(v, 1, lv, PartialQuickSort(lo:hi), Base.Sort.Forward)
-    end
-    isnan(v[end]) && throw(ArgumentError("quantiles are undefined in presence of NaNs"))
+    minp, maxp = extrema(p)
+    _quantilesort!(v, sorted, minp, maxp)
 
     for (i, j) in zip(eachindex(p), eachindex(q))
         @inbounds q[j] = _quantile(v,p[i])
@@ -671,21 +662,30 @@ end
 quantile!(v::AbstractVector, p::AbstractArray; sorted::Bool=false) =
     quantile!(similar(p,float(eltype(v))), v, p; sorted=sorted)
 
-function quantile!(v::AbstractVector, p::Real;
-                   sorted::Bool=false)
+quantile!(v::AbstractVector, p::Real; sorted::Bool=false) =
+    _quantile(_quantilesort!(v, sorted, p, p), p)
+
+function quantile!(v::AbstractVector, p::Tuple{Vararg{Real}}; sorted::Bool=false)
+    isempty(p) && return ()
+    minp, maxp = extrema(p)
+    _quantilesort!(v, sorted, minp, maxp)
+    return map(x->_quantile(v, x), p)
+end
+
+# Function to perform partial sort of v for quantiles in given range
+function _quantilesort!(v::AbstractArray, sorted::Bool, minp::Real, maxp::Real)
     isempty(v) && throw(ArgumentError("empty data vector"))
 
-    lv = length(v)
     if !sorted
-        lo = floor(Int,1+p*(lv-1))
-        hi = ceil(Int,1+p*(lv-1))
+        lv = length(v)
+        lo = floor(Int,1+minp*(lv-1))
+        hi = ceil(Int,1+maxp*(lv-1))
 
         # only need to perform partial sort
         sort!(v, 1, lv, PartialQuickSort(lo:hi), Base.Sort.Forward)
     end
     isnan(v[end]) && throw(ArgumentError("quantiles are undefined in presence of NaNs"))
-
-    return _quantile(v,p)
+    return v
 end
 
 # Core quantile lookup function: assumes `v` sorted

diff --git a/test/statistics.jl b/test/statistics.jl
@@ -326,6 +326,10 @@ end
 @test quantile([Inf,Inf],0.5) == Inf
 @test quantile([-Inf,1],0.5) == -Inf
 @test quantile([0,1],1e-18) == 1e-18
+@test quantile([1, 2, 3, 4],[]) == []
+@test quantile([1, 2, 3, 4], (0.5,)) == (2.5,)
+@test quantile([4, 9, 1, 5, 7, 8, 2, 3, 5, 17, 11], (0.1, 0.2, 0.4, 0.9)) == (2.0, 3.0, 5.0, 11.0)
+@test quantile([1, 2, 3, 4], ()) == ()
 
 # StatsBase issue 164
 y = [0.40003674665581906,0.4085630862624367,0.41662034698690303,0.41662034698690303,0.42189053966652057,0.42189053966652057,0.42553514344518345,0.43985732442991354]