用前非NA值起始的递增序列填充向量NA值——求更优实现

我有一个由整数和NA值组成的向量，需要将其中连续的NA值替换为以前一个非NA值为起始的递增整数序列。

现有实现方案 #

方案1：逻辑直观但速度较慢 #

使用Reduce函数实现，代码简洁但性能表现差：

v = c(NA, NA, 5, NA, NA, 2, 8, NA)
Reduce(\(i,j) if(is.na(j)) i+1 else j, v, accumulate = T)
# 输出结果：[1] NA NA  5  6  7  2  8  9

方案2：速度较快但写法繁琐 #

基于rle和bit::vecseq实现，基准测试显示比方案1快约10倍，但代码可读性较差：

r = rle(is.na(v)); w = which(r$values[-1]); s = cumsum(r$lengths); v2 = v
v2[bit::vecseq(s[w]+1, s[w+1])] = bit::vecseq(v[s[w]]+1, v[s[w]]+r$lengths[w+1]); v2
# 输出结果：[1] NA NA  5  6  7  2  8  9

是否存在更快或更简洁的实现方案？

补充实现与基准测试 #

以下是补充的多种实现方案，以及基于microbenchmark的性能对比：

测试代码 #

set.seed(0); v = sample(2e4); v[sample(2e4, 1e4)] = NA
vnum = as.numeric(v)

# last observation incremented forward
loif = \(x){r = rle(is.na(x)); w = which(r$values[-1]); s = cumsum(r$lengths)
  x[bit::vecseq(s[w]+1, s[w+1])] = bit::vecseq(v[s[w]]+1, v[s[w]]+r$lengths[w+1]); x}

fill_incr <- function(v){
  idx = which(complete.cases(v))
  le = length(v)

  d <- diff(c(idx, le + 1))
  v[idx[1]:le] <- sequence(d, v[complete.cases(v)])
  v
}

library(collapse)
fill_incr_collapse <- function(v){
  idx = whichNA(v, invert = TRUE)
  le = length(v)

  d <- diff(c(idx, le + 1))
  v[idx[1]:le] <- sequence(d, na_rm(v))
  v
}

# 用单次!is.na替换两次complete.cases调用
fill_incr2 = \(x){nn = !is.na(x); w = which(nn); l = length(x); d = diff(c(w, l+1))
  x[w[1]:l] = sequence(d, x[nn]); x}

cpp11::cpp_source(,'#include "cpp11.hpp"
using namespace cpp11;
[[cpp11::register]]
doubles loif_cpp(doubles xs) {
  bool hadnonna = false;
  int n = xs.size();
  writable::doubles out(n) ;
  for (int i = 0; i < n; ++i) {
    if (hadnonna && ISNA(xs[i])) {
      out[i] = out[i-1] + 1;
    } else {
      out[i] = xs[i];
      if (!ISNA(xs[i])) hadnonna = true;
    }
  }
  return out;
}')

cppFunction('NumericVector na_locf_numeric(NumericVector x) {
  int n = x.size();
  LogicalVector ina = is_na(x);

  for(int i = 1; i<n; i++) {
    if(ina[i] == TRUE) {
      x[i] = x[i-1] + 1;
    }
  }
  return x;
}')

b = microbenchmark::microbenchmark(times=100,
  Reduce(\(i,j)if(is.na(j))i+1 else j,v,accumulate=T),
  loif(v),
  fill_incr(v),
  fill_incr2(v),
  fill_incr_collapse(v),
  loif_cpp(vnum),
  na_locf_numeric(vnum))

o = sort(tapply(b$time, gsub("  +"," ",b$expr), median))
writeLines(sprintf("%.2f %s", o/min(o), names(o)))

基准测试结果（100次运行中位数时间，相对最快方案） #

1.00 na_locf_numeric(vnum)
2.26 loif_cpp(vnum)
9.80 fill_incr_collapse(v)
11.63 fill_incr2(v)
15.63 fill_incr(v)
20.41 loif(v)
211.83 Reduce(function(i, j) if (is.na(j)) i + 1 else j, v, accumulate = T)

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