使用Rcpp实现dplyr::group_by的分组索引功能

嘿，我来帮你搞定这个用Rcpp实现类似dplyr分组索引的需求！其实不用纠结std::unordered_multimap，咱们换个更直观高效的思路——直接用C++的映射容器来收集每个分组对应的原索引，完全贴合dplyr的输出格式。

方案1：按分组值自然排序输出（和默认dplyr行为一致） #

这个版本会自动按分组值的升序返回索引列表，和dplyr::group_by的默认分组顺序匹配：

#include <Rcpp.h>
#include <map>
#include <vector>

using namespace Rcpp;

// [[Rcpp::export]]
List group_indices(IntegerVector x) {
  // 用map存分组值到索引列表的映射，自动按key排序
  std::map<int, std::vector<int>> groups;
  
  // 遍历输入向量，把每个元素的原索引（从0开始）加到对应分组里
  for (int i = 0; i < x.size(); ++i) {
    groups[x[i]].push_back(i);
  }
  
  // 把map里的索引列表转成R的List返回
  List result;
  for (auto& group_pair : groups) {
    result.push_back(wrap(group_pair.second));
  }
  
  return result;
}

方案2：按分组首次出现的顺序输出（对应dplyr的.groups = "keep"） #

如果你需要和分组首次出现的顺序一致，而不是按值排序，就用这个版本：

#include <Rcpp.h>
#include <unordered_map>
#include <vector>

using namespace Rcpp;

// [[Rcpp::export]]
List group_indices_ordered(IntegerVector x) {
  std::unordered_map<int, std::vector<int>> groups;
  std::vector<int> group_order; // 记录分组第一次出现的顺序
  
  for (int i = 0; i < x.size(); ++i) {
    int current_val = x[i];
    // 如果是新分组，先记录顺序
    if (groups.find(current_val) == groups.end()) {
      group_order.push_back(current_val);
    }
    // 把当前索引加到对应分组
    groups[current_val].push_back(i);
  }
  
  // 按首次出现顺序构建结果List
  List result;
  for (int val : group_order) {
    result.push_back(wrap(groups[val]));
  }
  
  return result;
}

测试验证 #

在R里跑你的示例数据试试：

x <- sample(1:3, 10, TRUE)
x # 示例输出：[1] 3 3 3 1 3 1 3 2 3 2

# 用方案1测试
group_indices(x)
# 输出和dplyr的indices完全一致：
# [[1]]
# [1] 3 5
# 
# [[2]]
# [1] 7 9
# 
# [[3]]
# [1] 0 1 2 4 6 8

为什么不用unordered_multimap？ #

unordered_multimap是键可重复的无序容器，要提取同一分组的所有索引需要额外遍历查找，反而不如直接用map或unordered_map在遍历输入时就收集索引来得高效，代码也更简洁易读。

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