在Rcpp/RcppArmadillo中高效实现矩阵负下标索引提取子矩阵

核心思路 #

R中A[-j, -j]的本质是删除指定索引的行和列，RcppArmadillo虽无直接语法，但可通过构造保留索引或逻辑筛选实现高效操作，以下是两种实用方案：

方法1：逻辑索引筛选 #

利用Armadillo的逻辑向量标记保留位置，再提取对应子矩阵：

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

// [[Rcpp::export]]
arma::mat submat_exclude(arma::mat A, arma::uvec j) {
  int n = A.n_rows;
  // 初始化逻辑向量，默认保留所有行/列
  arma::vec keep = arma::vec(n, arma::fill::ones);
  // 转换为0-based索引后，标记需删除的位置
  keep.elem(j - 1) = 0;
  // 提取需保留的索引
  arma::uvec pos = arma::find(keep == 1);
  // 返回子矩阵
  return A(pos, pos);
}

方法2：直接构造保留索引（性能更优） #

跳过需删除的索引，直接生成保留序列，避免逻辑向量的额外开销：

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

// [[Rcpp::export]]
arma::mat submat_exclude_fast(arma::mat A, arma::uvec j) {
  int n = A.n_rows;
  // 转换为0-based并排序，方便遍历跳过
  arma::uvec j_sorted = arma::sort(j - 1);
  arma::uvec pos(n - j.n_elem);
  int idx = 0;
  int j_ptr = 0;
  
  for (int i = 0; i < n; ++i) {
    if (j_ptr < j_sorted.n_elem && i == j_sorted(j_ptr)) {
      j_ptr++;
    } else {
      pos(idx++) = i;
    }
  }
  
  return A(pos, pos);
}

关键注意事项 #

• Armadillo采用0-based索引，R传入的j是1-based，必须先减1转换，否则会触发索引越界错误。
• 第二种方法在矩阵规模大、需删除的索引数量多时，性能优于逻辑索引方案。

R端验证代码 #

A <- matrix(1:16, 4, 4)
j <- c(2, 3)
# R原生实现
A[-j, -j]
# 调用Rcpp函数验证
submat_exclude(A, j)
submat_exclude_fast(A, j)

内容的提问来源于stack exchange，提问作者Søren Højgaard