improve r_split

jl-pkgs · Sep 22, 2024 · e713dcb · e713dcb
1 parent ee93dae
commit e713dcb
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/src/r_base.jl b/src/r_base.jl
@@ -1,3 +1,4 @@
+import DataStructures: OrderedDict
 # using LoopVectorization: @turbo
 
 # function r_in(x::AbstractVector, y::AbstractVector)::BitVector
@@ -59,20 +60,19 @@ end
 #   DataFrame(vals, names)
 # end
 
-
 function r_split(lst::AbstractVector, by::AbstractVector)
   grps = unique(by) |> sort
-  res = []
+  res = OrderedDict()
   for grp in grps
     inds = by .== grp
+    res[grp] = lst[inds]
     # ans = Dict(grp => lst[inds])
-    ans = lst[inds]    
-    push!(res, ans)
+    # ans = lst[inds]
+    # push!(res, ans)
   end
   res
 end
 
-
 r_summary(x; digits=2) = nothing
 
 function r_summary(x::AbstractArray{<:Real}; digits=2)

diff --git a/test/test-r_base.jl b/test/test-r_base.jl
@@ -65,8 +65,9 @@ end
 
 
 @testset "r_split" begin
-  @test r_split(1:4, [1, 1, 1, 2]) == [[1, 2, 3], [4]]
-  @test r_split(1:4, [1, 1, 2, 2]) == [[1, 2], [3, 4]]
+  r = r_split(1:4, [1, 1, 1, 2])
+  @test r[1] == [1, 2, 3]
+  @test r[2] == [4]
 end
 
 @testset "writelines" begin