Stan equivalent to R's which()

Algorithms
13

Just came across this when answering a different post. Here are slightly modified versions of Bob’s suggestions above that I just tested when responding to How to find the location of a value in a vector. I changed it to work with a vector input (instead of an int array) and fixed a typo where match_positions[pos] = x[i] should be match_positions[pos] = i.

// Find the number of elements in the vector x that equal real number y
int num_matches(vector x, real y) {
  int n = 0;
  for (i in 1:rows(x))
    if (x[i] == y)
      n += 1;
  return n;
}
  
// Find the indexes of the elements in the vector x that equal real number y
int[] which_equal(vector x, real y) {
  int match_positions[num_matches(x, y)];
  int pos = 1;
  for (i in 1:size(x)) {
    if (x[i] == y) {
      match_positions[pos] = i;
      pos += 1;
    }
  }
  return match_positions;
}

// example usage later in stan program
int ids[num_matches(x, y)] = which_equal(x, y);
4 Likes