How to combine thrust comparisons using placeholders?

Question

I have a device_vector of float values A of size N. I also have a float value V for comparison. Depending on an input value I need to extract indices of A for which the values are > OR < OR == V.

I use the following code but it seems cumbersome. Is there a more concise way to do it?

void detect_indices_lesser_greater_equal_to_value(thrust::device_vector<float> S, float value, 
                        int criterion, thrust::device_vector<int>& indices)
{

int N=S.size();

int size=N;


if(criterion==0) // criterion =0 => equal
{
thrust::device_vector<int>::iterator end = thrust::copy_if(thrust::device,thrust::make_counting_iterator(0),
                                                             thrust::make_counting_iterator(N),
                                                             S.begin(),
                                                             indices.begin(), 
                                                             thrust::placeholders::_1 == value);
size = end-indices.begin();
}


if(criterion==1) // criterion =1 => less
{
thrust::device_vector<int>::iterator end = thrust::copy_if(thrust::device,thrust::make_counting_iterator(0),
                                                             thrust::make_counting_iterator(N),
                                                             S.begin(),
                                                             indices.begin(), 
                                                             thrust::placeholders::_1 < value);
size = end-indices.begin();
}

if(criterion==2) // criterion =2 => greater
{
thrust::device_vector<int>::iterator end = thrust::copy_if(thrust::device,thrust::make_counting_iterator(0),
                                                             thrust::make_counting_iterator(N),
                                                             S.begin(),
                                                             indices.begin(), 
                                                             thrust::placeholders::_1 > value);
size = end-indices.begin();
}

indices.resize(size);

}

Answer 1

This can be done with two thrust::partition operations. Partitioning is pretty simple: everything that results in a true predicate is moved to the left side of the input vector. Everything else is moved to the right. Here's a simple example:

$ cat t22.cu
#include <thrust/partition.h>
#include <thrust/copy.h>
#include <thrust/device_vector.h>

typedef float mt;

using namespace thrust::placeholders;
int main(){

  const mt pval = 4;
  mt data[] = {1,3,7,4,5,2,4,3,9};
  const int ds = sizeof(data)/sizeof(data[0]);
  thrust::device_vector<mt> d(data, data+ds);
  auto end1 = thrust::partition(d.begin(), d.end(), _1<pval);
  auto end2 = thrust::partition(end1, d.end(), _1==pval);
  std::cout << "less than pval:" << std::endl;
  thrust::copy(d.begin(), end1, std::ostream_iterator<mt>(std::cout,","));
  std::cout << std::endl << "equal to pval:" << std::endl;
  thrust::copy(end1, end2, std::ostream_iterator<mt>(std::cout,","));
  std::cout << std::endl << "greater than pval:" << std::endl;
  thrust::copy(end2, d.end(), std::ostream_iterator<mt>(std::cout,","));
  std::cout << std::endl;
}
$ nvcc -o t22 t22.cu
$ ./t22
less than pval:
1,3,2,3,
equal to pval:
4,4,
greater than pval:
7,5,9,
$

If you require that the ordering in the 3 resultant sub-vectors be the same as the original input ordering, you could use the thrust::stable_partition variant.

(Note that in your question you refer to float quantities, but your example code uses <int> iterators. However the above code can work with either by modifying the typedef).