Mean by groups Rcpp

Question

Here is an attempt to translate a R function to Rcpp. The goal is to group every n observations together (of a numeric vector) and calculate the mean.

On big'ish vectors the R function is almost two order of magnitude slower (when n is small). But I wonder:

1. If there are even more speed to be gained.
2. There are any elements of bad or (horrible) practice.

mean_every_n <- function(x, n=5L) {
  grp <- (seq_along(x) - 1 ) %/% n
  ave(x, grp, FUN = \(y) mean(y, na.rm = TRUE))
}

NumericVector mean_every_n(const NumericVector x, int n) {
  const int nx = x.size();
  double csum = 0; // cumulative sum
  double nnas = 0; // number of NAs
  NumericVector out(nx);
  for (int i = 0; i < nx; i++) {
    if (NumericVector::is_na(x[i])) nnas++; else csum += x[i];
    bool reached_end = (i == (nx-1));
    if ((i+1) % n == 0 || reached_end) {
      if (reached_end && (i+1) % n > 0) n = (i+1) % n;
      if (nnas < n) {
        double res = csum / (n - nnas);
        for (int j = 0; j < n; j++) out[i - j] = res;
      } else {
        for (int j = 0; j < n; j++) out[i - j] = NA_REAL;
      }
      csum = 0, nnas = 0;
    }
  }
  return out;
}

Inspired by a SO question.

Answer 1

Pass by const reference

You are passing the NumericVector x by value. This will cause a copy to be made. It is better to pass it by const reference.

Simplify the code

Your C++ code is a very low level implementation of what you want, with everything in one big function. You can greatly improve it by splitting it up into multiple functions. For example, one that splits your input into chunks of size n, and another that calculates the average of a chunk. Consider how the R function makes use of seq_along(), ave() and mean(). You can implement such helper functions yourself, but that brings me to:

Make use of the latest version of C++

The ranges library in C++23 provides a lot of what you need. You could rewrite your code to something like:

double mean(const NumericVector& x) {
    return std::ranges::fold_left_first(x, std::plus<double>()) / x.size();
}

NumericVector mean_every_n(const NumericVector& x, int n) {
    return std::views::chunk(x, n)
         | std::views::transform(mean)
         | std::ranges::to<NumericVector>();
}

Note that this outputs only one value for each mean, and I've also omitted the handling of NaNs here. Adding that is left as an exercise for the reader. Again, consider creating a helper function for that, for example a version of plus() that returns a NaN if both arguments are a NaN, the non-NaN value if only one argument is not NaN, and the sum if both arguments are not NaN.

More speed

A good C++ compiler, with optimizations enabled, will probably generate pretty good single-threaded code. However, if you have very large inputs, then you could use multiple threads to process the input in parallel. This can be done either manually, by using OpenMP, or possibly using some of the C++ standard library's parallel algorithms (see for example the verison of std::transform() that takes an execution policy as a parameter).