Search code examples
c++eigeneigen3

Eigen3 select rows out based on column conditions


I have a matrix in eigen which is 2 dimensions, such as:

122 443 544 456 0.9
324 435 5465 645 0.8
32 434 545 546 0.778
435 546 6565 656 0.6878
546 6565 656 3453 54 0.7788
5456 546 545 6565 3434 0.244
435 5456 656 656 6565 0.445
.....

I want select all rows (or get it's row index out) when the last column value bigger than 0.3.

I know I can do this by iteration all rows and judge the last element, but I have maybe 10000 rows, to do this, iteration will be very slow.

Does there any better way to do this?


Solution

  • The selection of the relevant rows can be done in a single line, by storing the result of the comparison of all the elements in the last column into a boolean array, which can be cast into a VectorXi.

    VectorXi is_selected = (mat.col(last_col).array() > 0.3).cast<int>();
    

    This information can then be used to prepare a new matrix which contains only the selected rows. An entire code using this method is shown below.

    #include <Eigen/Dense>
    #include <iostream>    
    using namespace Eigen;
    
    int main() {
      const int nr = 10;
      const int nc = 5;
      MatrixXd mat = MatrixXd::Random(nr,nc);
      std::cout << "original:\n" << mat << std::endl;
      int last_col = mat.cols() - 1;
    
      VectorXi is_selected = (mat.col(last_col).array() > 0.3).cast<int>();
    
      MatrixXd mat_sel(is_selected.sum(), mat.cols());
      int rownew = 0;
      for (int i = 0; i < mat.rows(); ++i) {
        if (is_selected[i]) {       
           mat_sel.row(rownew) = mat.row(i);
           rownew++;
        }
      }
      std::cout << "selected:\n" << mat_sel << std::endl;
    }
    

    Demo: https://godbolt.org/z/f0_fC0

    Edit: Using new feature (Eigen 3.4 or 3.3.90 development branch)

    The development branch of Eigen provides a new overload of the MatrixX constructor that allows for the direct subsetting of a given matrix.

    MatrixXd mat_sel = mat(keep_rows, keep_cols); 
    

    The columns and rows that should be kept are stored in an Eigen::VectorXi or in a std::vector<int>:

    #include <Eigen/Dense>
    #include <iostream>
    #include <vector>
    using namespace Eigen;
    
    int main() {
      MatrixXd mat = MatrixXd::Random(10,5);
      std::cout << "original:\n" << mat << std::endl;
      std::vector<int> keep_rows;  
      for (int i = 0; i < mat.rows(); ++i) {
        if (mat(i,mat.cols() - 1) > 0.3) {
           keep_rows.push_back(i);
         }     
      }
      VectorXi keep_cols = VectorXi::LinSpaced(mat.cols(), 0, mat.cols());
      MatrixXd mat_sel = mat(keep_rows, keep_cols);          
      std::cout << "selected:\n" << mat_sel << std::endl; 
    }
    

    Demo: https://godbolt.org/z/Ag7g7f