R 查找一个向量的匹配元素在另一个向量中的位置

发布于03月16日

我有两个非常长的向量:

a <- sample(1e+08L, size = 1e+09L, replace = TRUE)
b <- sample(1e+08L, size = 1e+09L, replace = TRUE)

我想生成一个长度为length(a)的整数向量r，这样r[i]就是b中a[i]的索引.

我试了pmatch(a, b)，但速度太慢了.有没有更有效的方法？

以下是一个小示例所需的输出:

a <- c(1, 3, 5, 7, 8)
b <- c(3, 1, 7, 8, 5)
f(a, b)
## [1] 2 1 5 3 4

推荐答案

一种方法是使用Rcpp包，并在Rcpp内执行自定义pmatch函数，即

library(Rcpp)
library(microbenchmark)

cppFunction('
  IntegerVector cpp_pmatch(const NumericVector& b, const NumericVector& a) {
    IntegerVector result(b.size());
    
    for (int i = 0; i < b.size(); ++i) {
      double b_value = b[i];
      bool found = false;

      for (int j = 0; j < a.size(); ++j) {
        double a_value = a[j];

        if (a_value == b_value || (int)a_value == (int)b_value) {
          result[i] = j + 1;
          found = true;
          break;
        }
      }

      if (!found) {
        result[i] = NA_INTEGER;
      }
    }
    
    return result;
  }
')

这将产生以下结果:

cpp_pmatch(a, b)
[1] 2 1 5 3 4
pmatch(a, b)
[1] 2 1 5 3 4

Benchmark:个

benchmark_result <- microbenchmark(
      cpp_pmatch = cpp_pmatch(a, b),
      pmatch = pmatch(a, b),
      times = 1000L
    )

> Unit: nanoseconds
       expr  min   lq   mean median   uq   max neval
 cpp_pmatch  900 1000 1166.4   1000 1100 14500  1000
     pmatch 7400 7500 8229.5   7600 7800 40400  1000