julia数组按组选择前x行

Question

使用 julia，我想选择每组数组的前 x 行。

在下面的示例中，我希望前两行第二列等于 1.0，然后是前两行第二列等于 2.0，依此类推。

XX = [repeat([1.0], 6) vcat(repeat([1.0], 3), repeat([2.0], 3))]
XX2 = [repeat([2.0], 6) vcat(repeat([3.0], 3), repeat([4.0], 3))]
beg = [XX;XX2]

> 12×2 Matrix{Float64}:
>  1.0  1.0
>  1.0  1.0
>  1.0  1.0
>  1.0  2.0
>  1.0  2.0
>  1.0  2.0
>  2.0  3.0
>  2.0  3.0
>  2.0  3.0
>  2.0  4.0
>  2.0  4.0
>  2.0  4.0

最终的数组如下所示：

8×2 Matrix{Float64}:
 1.0  1.0
 1.0  1.0
 1.0  2.0
 1.0  2.0
 2.0  3.0
 2.0  3.0
 2.0  4.0
 2.0  4.0

我使用以下代码，但我不确定是否有一种更简单的方法（一个函数）可以更有效地做到这一点？

x = []
for val in unique(beg[:,2])
    x = append!(x, findfirst(beg[:,2].==val))
end
idx = sort([x; x.+1])
final = beg[idx, :] 

Answer 1

假设您的数据：

• 已排序（即组正在形成连续的块）
• 保证每个组至少有两个元素

（您的代码假定两者）

然后您可以通过以下方式生成您想要的idx过滤器：

idx == [i for i in axes(beg, 1) if i < 3 || beg[i, 2] != beg[i-1, 2] || beg[i, 2] != beg[i-2, 2]]

如果您不能假设以上任何一种，请发表评论，我可以提供更通用的解决方案。

编辑

这是一个不使用任何外部包的示例：

julia> using Random

julia> XX = [repeat([1.0], 6) vcat(repeat([1.0], 3), repeat([2.0], 3))]
6×2 Matrix{Float64}:
 1.0  1.0
 1.0  1.0
 1.0  1.0
 1.0  2.0
 1.0  2.0
 1.0  2.0

julia> XX2 = [repeat([2.0], 7) vcat(repeat([3.0], 3), repeat([4.0], 3), 5.0)] # last group has length 1
7×2 Matrix{Float64}:
 2.0  3.0
 2.0  3.0
 2.0  3.0
 2.0  4.0
 2.0  4.0
 2.0  4.0
 2.0  5.0

julia> beg = [XX;XX2][randperm(13), :] # shuffle groups so they are not in order
13×2 Matrix{Float64}:
 2.0  3.0
 1.0  2.0
 2.0  4.0
 2.0  3.0
 2.0  4.0
 2.0  5.0
 2.0  3.0
 1.0  2.0
 1.0  2.0
 1.0  1.0
 1.0  1.0
 2.0  4.0
 1.0  1.0

julia> x = Dict{Float64, Vector{Int}}() # this will store indices per group
Dict{Float64, Vector{Int64}}()

julia> for (i, v) in enumerate(beg[:, 2]) # collect the indices
           push!(get!(x, v, Int[]), i)
       end

julia> x
Dict{Float64, Vector{Int64}} with 5 entries:
  5.0 => [6]
  4.0 => [3, 5, 12]
  2.0 => [2, 8, 9]
  3.0 => [1, 4, 7]
  1.0 => [10, 11, 13]

julia> idx = sort!(mapreduce(x -> first(x, 2), vcat, values(x))) # get first two indices per group in ascending order
9-element Vector{Int64}:
  1
  2
  3
  4
  5
  6
  8
 10
 11