Julia中按分组选取Array数组每组前x行的高效实现方法

你可以根据场景选择以下两种实现方式：

方法1：使用DataFrames包（简洁易读，优先推荐） #

这是常规数据处理场景下最省心的写法，逻辑直观，后续修改需求也很方便：

using DataFrames

# 把原矩阵转为DataFrame，自动命名列
df = DataFrame(beg, :auto)
# 按第二列分组，每组取前2行，直接调整数字即可修改每组取数行数
res_df = combine(groupby(df, :x2), first => 2)
# 转回矩阵即得到目标结果
final = Matrix(res_df[:, [:x1, :x2]])

方法2：纯Base Julia实现（无额外依赖） #

如果不想引入第三方包，可以用下面的通用写法，支持任意调整每组取数行数，性能也优于你当前的实现：

group_col = beg[:, 2]
take_n = 2  # 每组取前n行，按需修改即可
res_idx = Int[]

for g in unique(group_col)
    # 取当前组的前n个索引，直接追加到结果索引数组
    append!(res_idx, findall(==(g), group_col)[1:take_n])
end

final = beg[res_idx, :]

这个写法不需要额外排序，因为组的顺序和组内行的顺序都和原数组完全一致，避免了不必要的排序开销。

内容的提问来源于stack exchange，提问作者djourd1