From e713dcb655719aff38fb3450b6f1e9cb8865d5a1 Mon Sep 17 00:00:00 2001
From: Dongdong Kong <kongdd@users.noreply.github.com>
Date: Sun, 22 Sep 2024 19:18:58 +0800
Subject: [PATCH] improve `r_split`

---
 src/r_base.jl       | 10 +++++-----
 test/test-r_base.jl |  5 +++--
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/r_base.jl b/src/r_base.jl
index 71f2d94..c91a961 100644
--- 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
index 23463b1..15c30f0 100644
--- 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