Allow for weights in cut(x,[ w,] ngroups)

JuliaData · Sep 20, 2019 · 930f0be · 930f0be
1 parent 0ede6dc
commit 930f0be
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -10,6 +10,7 @@ JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 Missings = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 
 [compat]

diff --git a/src/extras.jl b/src/extras.jl
@@ -1,4 +1,4 @@
-using Compat.Statistics
+import StatsBase
 
 function fill_refs!(refs::AbstractArray, X::AbstractArray,
                     breaks::AbstractVector, extend::Bool, allow_missing::Bool)
@@ -166,7 +166,7 @@ function cut(x::AbstractArray{T, N}, breaks::AbstractVector;
 end
 
 """
-    cut(x::AbstractArray, ngroups::Integer;
+    cut(x::AbstractArray, [w::AbstractWeights, ]ngroups::Integer;
         labels::Union{AbstractVector{<:AbstractString},Function})
 
 Cut a numeric array into `ngroups` quantiles, determined using
@@ -175,3 +175,7 @@ Cut a numeric array into `ngroups` quantiles, determined using
 cut(x::AbstractArray, ngroups::Integer;
     labels::Union{AbstractVector{<:AbstractString},Function}=default_formatter) =
     cut(x, Statistics.quantile(x, (1:ngroups-1)/ngroups); extend=true, labels=labels)
+
+cut(x::AbstractArray, w::AbstractWeights, ngroups::Integer;
+  labels::Union{AbstractVector{<:AbstractString},Function}=default_formatter) =
+  cut(x, StatsBase.quantile(x, w, (1:ngroups-1)/ngroups); extend=true, labels=labels)