Rcpp CV index

library(Rcpp)
library(RcppArmadillo)

sourceCpp(code = '
// [[Rcpp::depends(RcppArmadillo)]]
#include <RcppArmadillo.h>
using namespace Rcpp;
using namespace arma;

// [[Rcpp::export]]
uvec cvfold_index_f(const uword n, const uword fold)
{
  uword fold_n = n / fold, rem = n - fold_n * fold;
  uvec tmp_vec = linspace<uvec>(0, fold - 1, fold),
    index = vectorise(repmat(tmp_vec, fold_n, 1));
  if (rem > 0)
    index = join_cols(index, index.head(rem));

  // random permutation
  int j;
  for (int i = 0; i < n; i++)
  {
    j = as_scalar(randi<uvec>(1, distr_param(i, n - 1)));
    index.swap_rows(i, j);
  }
  return index;
}')
cvfold_index_f(11, 3) # 2 2 1 2 0 1 0 0 0 1 1