如何使用RcppArmadillo查找向量中与另一个向量相同的元素的索引？

Question

我现在卡住了，想要使用Rcpp Armadillo找到向量x中的元素在另一个向量vals中的索引。这两个向量x和vals的类型都是arma::uvec。

在R中，这很简单：

x <- c(1,1,1,4,2,4,4)
vals <- c(1,4)
which(v %in% vals)

我已经查看了Armadillo文档并尝试使用find()，但由于vals是一个向量，所以它无法正常工作。我还尝试过intersect()，但它只返回第一个唯一索引。

使用Armadillo，有什么好的/高效的方法可以实现这个目标吗？我是否必须使用find()来迭代遍历vals中的元素？

- Econ21

Answer 1

一个快速而简单的方法：

Rcpp::cppFunction("
  arma::uvec ind(arma::uvec x, arma::uvec y){
   arma::vec a(x.size(), arma::fill::zeros);
   for (auto i:y) a = a +  (x==i);
   return arma::find(a) + 1;
  }
 ", 'RcppArmadillo')

c(ind(v, vals))
[1] 1 2 3 4 6 7

- onyambu

---

甚至一點也不骯髒！我只能想到一種通過hash的方法，這樣更簡潔！ - Dirk Eddelbuettel

谢谢，比我的临时解决方案优雅得多。i:y是指直接遍历y中的元素，而不是通过“索引”（如for (i = 0; i < y.size(); i++) a = a + (x==y(i))）进行遍历吗？ - Econ21

@Econ21 是的，你说得没错。 - Onyambu

Answer 2

为了完整起见，我在此期间想出了这个解决方案：

arma::uvec getIndex(arma::uvec x, arma::uvec y) {
  
  int i, j, k = 0, n = y.size();
  arma::uvec tmp(n);
  
  for (i = 0; i < n; i++) {
    arma::uvec tmpID = arma::find(x == y(i));
    tmp(i) = tmpID.size();
  }
  
  arma::uvec out(sum(tmp));
  
  for (i = 0; i < n; i++) {
    arma::uvec id = arma::find(x == y(i));
    for (j = 0; j < id.size(); j++) {
      out(j+k) = id(j);
    }
    k += tmp(i);
  }
  return out;
}

- Econ21