diff --git a/src/Statistics.jl b/src/Statistics.jl index 91cf56fb..e278549e 100644 --- a/src/Statistics.jl +++ b/src/Statistics.jl @@ -937,8 +937,7 @@ function quantile!(q::AbstractArray, v::AbstractVector, p::AbstractArray; end isempty(q) && return q - minp, maxp = extrema(p) - _quantilesort!(v, sorted, minp, maxp) + _quantilesort!(v, sorted, p) for (i, j) in zip(eachindex(p), eachindex(q)) @inbounds q[j] = _quantile(v,p[i], alpha=alpha, beta=beta) @@ -949,27 +948,27 @@ end function quantile!(v::AbstractVector, p::Union{AbstractArray, Tuple{Vararg{Real}}}; sorted::Bool=false, alpha::Real=1., beta::Real=alpha) if !isempty(p) - minp, maxp = extrema(p) - _quantilesort!(v, sorted, minp, maxp) + _quantilesort!(v, sorted, p isa Tuple ? collect(p) : p) end return map(x->_quantile(v, x, alpha=alpha, beta=beta), p) end quantile!(v::AbstractVector, p::Real; sorted::Bool=false, alpha::Real=1., beta::Real=alpha) = - _quantile(_quantilesort!(v, sorted, p, p), p, alpha=alpha, beta=beta) + _quantile(_quantilesort!(v, sorted, [p]), p, alpha=alpha, beta=beta) -# Function to perform partial sort of v for quantiles in given range -function _quantilesort!(v::AbstractArray, sorted::Bool, minp::Real, maxp::Real) +function _quantilesort!(v::AbstractArray, sorted::Bool, p::AbstractArray) isempty(v) && throw(ArgumentError("empty data vector")) require_one_based_indexing(v) if !sorted - lv = length(v) - lo = floor(Int,minp*(lv)) - hi = ceil(Int,1+maxp*(lv)) - - # only need to perform partial sort - sort!(v, 1, lv, Base.Sort.PartialQuickSort(lo:hi), Base.Sort.Forward) + start = 1 + for pv in sort(p) + lv = length(v) + lo = floor(Int,pv*(lv)) + hi = ceil(Int,1+pv*(lv)) + sort!(v, start, lv, Base.Sort.PartialQuickSort(lo:hi), Base.Sort.Forward) + start = hi + 1 + end end if (sorted && (ismissing(v[end]) || (v[end] isa Number && isnan(v[end])))) || any(x -> ismissing(x) || (x isa Number && isnan(x)), v) @@ -977,7 +976,6 @@ function _quantilesort!(v::AbstractArray, sorted::Bool, minp::Real, maxp::Real) end return v end - # Core quantile lookup function: assumes `v` sorted @inline function _quantile(v::AbstractVector, p::Real; alpha::Real=1.0, beta::Real=alpha) 0 <= p <= 1 || throw(ArgumentError("input probability out of [0,1] range")) diff --git a/test/runtests.jl b/test/runtests.jl index 2e061ee4..fed79f14 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -601,6 +601,14 @@ end @test quantile(skipmissing([1, missing, 2]), 0.5) === 1.5 @test quantile([1], 0.5) === 1.0 + # randomized partialsort correctness test + Random.seed!(1234) + for i in 1:200, j in 1:20 + x = rand(2000) + p = rand(j) + @test quantile(x, p) == [quantile(x, v) for v in p] + end + # make sure that type inference works correctly in normal cases for T in [Int, BigInt, Float64, Float16, BigFloat, Rational{Int}, Rational{BigInt}] for S in [Float64, Float16, BigFloat, Rational{Int}, Rational{BigInt}]