R:大向量的高效迭代子集和过滤

Question

我想更快地执行以下操作。

逻辑:我有一个包含 4 个元素 1、2、3、4 的向量 big。我还有一个相同长度的阈值向量 1.1, 3.1, 4.1, 5.1。我希望每个元素找到第一个 next 元素的索引高于相应的阈值。在这种情况下，我的预期输出是

2、3、NA、NA:

• 第一个元素在第一个元素(包括)之后高于阈值 1.1 的第一个元素位于索引 2(值为 2)。
• 高于第二个阈值 3.1 的第一个元素的值为 4，并且是 在索引 2 处的当前元素之后的第三个元素(包括在内)。

---

基础实现

start <- Sys.time()
bigg <- rnorm(25000)
thresh <- bigg+0.5
result <- rep(NA, length(bigg))
for(i in 1:length(bigg)) {
  result[i] <- which(bigg[(i+1):length(bigg)]>thresh[i])[1] # the first next element that is higher than thresh
  if(i%%1000==0) print(paste0(i, " ", round(i/length(bigg),3)))
}
end <- Sys.time()
end-start
head(result)

基本上，取向量 x 的第一个元素在当前元素之后满足阈值条件。

我尝试使用 Rcpp

// [[Rcpp::export]]
int cppnextup_(NumericVector x, double thresh, bool is_up = true) {
  int n = x.size();
  //int idx = 0;
  int res = -1;
  for(int idx = 0; idx < n; ++idx) {
    if(x[idx]>thresh && is_up == true) {
      res = idx;
//Rcout << "The value of idx : " << idx <<" "<< x[idx]<<"\n";
      break;
    }
    if(x[idx]<thresh && is_up == false) {
      res = idx;
      //Rcout << "The value of idx : " << idx <<" "<< x[idx]<<"\n";
      break;
    }
  }
  return res;
}

基准测试:

# base --------------------------------------------------------------------

base_ <- function() {
  for(i in 1:length(bigg)) {
    result[i] <- which(bigg[(i+1):length(bigg)]>thresh[i])[1] # the first next element that is higher than thresh
    if(i%%1000==0) print(paste0(i, " ", round(i/length(bigg),3)))
  }
}

# cpp ----------------------------------------------------------------

result_cpp <- rep(NA, length(bigg))
cpp_ <- function() {
  for(i in 1:length(bigg)) {
    result_cpp[i] <- cppnextup_(bigg[(i+1):length(bigg)], thresh[i]) # the first next element that is higher than thresh
    if(i%%1000==0) print(paste0(i, " ", round(i/length(bigg),3)))
  }
}

#result_cpp <- ifelse(result_cpp==-1, NA, result_cpp)
#result_cpp <- result_cpp+1
#all.equal(result, result_cpp)
#[1] TRUE

# benchmark ---------------------------------------------------------------

microbenchmark::microbenchmark(base_(),
                               cpp_(), times=3)
Unit: milliseconds
    expr      min        lq      mean    median        uq       max neval
 base_() 2023.510 2030.3154 2078.7867 2037.1211 2106.4252 2175.7293     3
  cpp_()  661.277  665.3456  718.8851  669.4141  747.6891  825.9641     3

我的 Rcpp 实现将基本时间减少了 65%，有没有更好的(矢量化)方法？寻找任何后端，无论是 Rcpp、data.table、dtplyr 等

我的 dtplyr 尝试产生所有 NA 的:

library(dtplyr)
nx <- length(bigg)
df <- tibble(bigg, thresh)
bigg %>% lazy_dt() %>% mutate(res = which(bigg[row_number():nx]>thresh)[1])
Warning message:
In seq_len(.N):..nx :
  numerical expression has 25000 elements: only the first used

干杯

顺便说一句，我的真实向量有 8,406,600 个元素。

---

编辑:矢量化 Rcpp

我还有另一个更快的 Rcpp 函数，它依赖于第一个函数:

// [[Rcpp::export]]
NumericVector cppnextup(NumericVector x, double threshup, bool is_up = true) {
  int n = x.size();
  NumericVector up(n);
  if(is_up == true) {
    up = x + threshup;
  } else {
    up = x - threshup;
  }
//  Rcout << "The value of up : " << up[0] <<" "<< up[1] <<"\n";
  NumericVector result(n);
  int idx = 0;
  for(int i = 0; i < n; ++i) {
    double thisup = up[idx];
    NumericVector thisvect = x[Rcpp::Range((idx), (n-1))];
    
//Rcout <<idx<< " " << "thisvect : " << thisvect[0] <<" thisup: "<< thisup <<" buy " << buy << "\n";
    
    int resi = cppnextup_(thisvect, thisup, is_up = is_up);
    if(resi != 0) {
      result[idx] = resi+1;
    } else {
      result[idx] = resi;
    }
    
    //Rcout << "RESI: " << resi <<" "<< up[1] <<"\n";
    idx = idx + 1;
  }
  return result;
}

如你所见，它比前两个更快:

# cpp_vectorized ----------------------------------------------------------

cpp_vect <- function(bigg) {
  res_cppvect <- cppnextup(bigg, 0.5)
}

    # benchmark ---------------------------------------------------------------
    
    microbenchmark::microbenchmark(base_(),
                                   cpp_(), 
                                   cpp_vect(),
                                   times=3)
           expr       min        lq      mean    median        uq       max neval
        base_() 2014.7211 2016.8679 2068.9869 2019.0146 2096.1198 2173.2250     3
         cpp_()  663.0874  666.1540  718.5863  669.2207  746.3357  823.4507     3
     cpp_vect()  214.1745  221.2103  223.9532  228.2460  228.8426  229.4392     3

但是当我在参数中传递一个更大的向量时，它会卡住并且永远不会返回结果。

res <- cpp_vect(bigg=rnorm(1000000)) # freezes

欢迎任何帮助。

Answer 1

使用 mult = "first" 的 data.table 非等值连接效果很好。不过，它不会像优化的 Rcpp 函数那样快。

library(data.table)

bigg <- rnorm(25000)
thresh <- bigg+0.5
f1 <- function(bigg, thresh) {
  result <- rep(NA, length(bigg))
  for(i in 1:length(bigg)) {
    result[i] <- which(bigg[(i+1):length(bigg)]>thresh[i])[1] # the first next element that is higher than thresh
  }
  result
}

f2 <- function(bigg, thresh) {
  data.table(
    val = bigg,
    r = seq_along(bigg)
  )[
    data.table(
      val = thresh,
      r = seq_along(thresh)
    ),
    on = .(val > val, r > r),
    .(result = x.r - i.r),
    mult = "first"
  ]$result
}

microbenchmark::microbenchmark(f1 = f1(bigg, thresh),
                               f2 = f2(bigg, thresh),
                               times = 10,
                               check = "identical")
#> Unit: milliseconds
#>  expr      min       lq      mean    median       uq       max neval
#>    f1 2167.139 2199.801 2217.6945 2222.4937 2233.254 2250.1693    10
#>    f2  605.999  610.576  612.0431  611.1439  614.195  618.6248    10

bigg <- rnorm(1e6)
thresh <- bigg+0.5
system.time(f2(bigg, thresh))
#>    user  system elapsed 
#>  375.71    0.15  375.81