RcppArmadillo sorting matrix based on orders of corresponding vectors

sort-matrix-using-vectors.R

library(Rcpp)

# Create an R function for sorting a matrix
# based on an accompanying vectors

cpp_code <- (
  'using namespace Rcpp;

//#include <Rcpp.h>
#include <RcppArmadillo.h>
// [[Rcpp::depends(RcppArmadillo)]]

// [[Rcpp::export()]]
arma::mat sort_matrix(arma::mat Y, arma::colvec samp_unit_ids, arma::colvec strata_ids) {

  // Get distinct strata ids and their length, H
  arma::colvec distinct_strata_ids = unique(strata_ids);
  arma::uword H = distinct_strata_ids.n_elem;

  // First reorder inputs by sample unit IDs
  arma::uvec samp_unit_id_order = arma::stable_sort_index(samp_unit_ids, "ascend");
  Y = Y.rows(samp_unit_id_order);
  samp_unit_ids = samp_unit_ids.elem(samp_unit_id_order);
  strata_ids = strata_ids.elem(samp_unit_id_order);

  // Next reorder inputs by strata IDs
  arma::uvec strata_id_order = arma::stable_sort_index(strata_ids, "ascend");
  Y = Y.rows(strata_id_order);
  samp_unit_ids = samp_unit_ids.elem(strata_id_order);
  strata_ids = strata_ids.elem(strata_id_order);

  return Y;
}'
) |>
cppFunction(depends = "RcppArmadillo")

my_strata <- c(1,2,2,1)
my_samp_units <- c(2,2,1,1)
my_matrix <- matrix(c(1,1,3,3,4,4,2,2),
                    nrow = 4, ncol = 2,
                    byrow = TRUE)
sort_matrix(my_matrix, my_samp_units, my_strata)