How do I find the indices of elements in a vector which are also in another vector using RcppArmadillo?

Question

I am stuck trying to find the indices of elements in a vector x whose elements are also in another vector vals using Rcpp Armadillo. Both x and vals are of type arma::uvec.

In R, this would be straightforward:

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

I've scanned the Armadillo docs and find() was my obvious first try; but it didn't work, since vals is a vector. I've also tried intersect() but it returns only the first unique indices.

What would be a good/efficient way to do this using Armadillo? Do I have to iterate through the elements in vals using find()?

score 5 · Accepted Answer · answered Jul 05 '22 at 00:03

5

A quick dirty way:

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

answered Jul 05 '22 at 00:03

Onyambu

67,392
3
24
53

1

Not even all that dirty! I could only think of a way via a `hash`, this is shorter! – Dirk Eddelbuettel Jul 05 '22 at 00:53
Thanks, much more elegant than my interim solution. Does `i:y` mean you loop through the elements in `y` directly instead of through the "index" (as in `for (i = 0; i < y.size(); i++) a = a + (x==y(i))`? – Econ21 Jul 05 '22 at 13:24
1

@Econ21 yes you are right – Onyambu Jul 05 '22 at 14:57

score 0 · Answer 2 · answered Jul 05 '22 at 13:19

For completeness, I came up with this solution in the meantime:

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;
}

How do I find the indices of elements in a vector which are also in another vector using RcppArmadillo?

2 Answers2