Rcpp代码正确性咨询：矩阵列两两KS检验结果数量不符问题

问题概述 #

要对矩阵每一列两两执行KS检验，当矩阵有5列时，预期得到5*4/2=10个结果，但现有代码仅返回5个，核心问题出在结果向量的初始化长度和索引变量m的使用上。

代码问题定位 #

1. 结果向量长度错误：原代码中Rcpp::NumericVector results(n);将结果向量长度设为列数n，但实际需要存储的两两组合数是n*(n-1)/2，导致超出长度的结果无法被保存。
2. 索引变量起始值错误：Rcpp的向量采用0索引，但m初始值设为1，导致第一个结果存在results[1]，results[0]保留默认的0值，同时后续结果因向量长度不足被截断。

修复后的代码 #

library(Rcpp)
sourceCpp(code='
// [[Rcpp::depends(RcppArmadillo)]]
#include <RcppArmadillo.h> 

double KS(arma::colvec x, arma::colvec y) {
  int n = x.n_rows;
  arma::colvec w = join_cols(x, y);
  arma::uvec z = arma::sort_index(w);
  w.fill(-1); w.elem( find(z <= n-1) ).ones();
  return max(abs(cumsum(w)))/n;
}
// [[Rcpp::export]]
Rcpp::NumericVector K_S(arma::mat mt) {
  int n = mt.n_cols;
  // 计算正确的两两组合数量
  int total = n * (n - 1) / 2;
  Rcpp::NumericVector results(total);
  // 索引从0开始
  int m = 0;
  for (int i = 0; i < n-1; i++) {
    for (int j = i+1; j < n; j++){
      arma::colvec x=mt.col(i);
      arma::colvec y=mt.col(j);
      results[m] = KS(x, y);
      m ++;
    }
  }
  return results;
}
')

set.seed(1)
mt <- matrix(rnorm(400*5), ncol=5)
result <- K_S(t(mt))

验证结果 #

运行修复后的代码，result会返回10个数值，符合预期：

> result
 [1] 0.1050 0.0675 0.0475 0.0650 0.0575 0.0625 0.0550 0.0600 0.0525 0.0450

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