Extracting entries at multiple indices from std::vectors in c++, from zipped matrices

Alec Jacobson

July 20, 2010

weblog/

As a naturally follow up to my last post, I have extended my C++ implementation of matlab like array indexing to work for "zipped matrices". I often store matrices zipped into std::vectors, column-wise. Now I can access them like how I do in matlab. Suppose in matlab I had:
A = [10,11,12;13,14,15];
You could then suck out a re-ordered block of rows and columns of the matrix by calling:
B = A([1],[3,2]);
Revealing:
B =

    12    11
This is especially convenient for chopping up and reordering rows and columns in linear systems. I can do this same sort of indexing with C++, the caveat being that I have to pass the number of columns as a third argument. For now, I think I would rather do this than have it be a field of the SmartMatrix.
#include <vector>
#include <stdio.h>

template<typename T>
class SmartMatrix: public std::vector<T>{
  public:
    // act like operator[]
    T operator()(
        size_t _row,
        size_t _col,
        size_t number_of_columns){
      return (*this)[_row*number_of_columns+_col];
    }

    // act like matlab operator()
    SmartMatrix<T> operator()(
        std::vector<size_t>& row_positions,
        std::vector<size_t>& column_positions, 
        size_t number_of_columns){
      SmartMatrix<T> sub;
      sub.resize(row_positions.size()*column_positions.size());
      size_t sub_i = 0;
      for(
          std::vector<size_t>::iterator rpit = row_positions.begin();
          rpit != row_positions.end();
          rpit++,sub_i++){
        size_t sub_j = 0;
        for(
          std::vector<size_t>::iterator cpit = column_positions.begin();
          cpit != column_positions.end();
          cpit++,sub_j++){
          sub[sub_i*column_positions.size() + sub_j] = 
            (*this)[(*rpit)*number_of_columns + (*cpit)];
        }
      }
      return sub;
    }
};

int main(){
  SmartMatrix<int> A;
  size_t cols_in_A = 3;
  size_t rows_in_A = 2;
  A.push_back(10);
  A.push_back(11);
  A.push_back(12);
  A.push_back(13);
  A.push_back(14);
  A.push_back(15);
  A.push_back(16);

  printf("A=\n");
  for(int i=0;i<rows_in_A;i++){
    for(int j=0;j<cols_in_A;j++)
      printf("%d ",A[i*cols_in_A+j]);
    printf("\n");
  }
  printf("\nA[%d,%d]=%d\n",0,2,A(0,2,cols_in_A));

  // make some list of indices
  std::vector<size_t> sub_row_indices;
  sub_row_indices.push_back(0);
  std::vector<size_t> sub_column_indices;
  sub_column_indices.push_back(2);
  sub_column_indices.push_back(1);

  // use operator() to extract entries at those indices
  SmartMatrix<int> B = A(sub_row_indices,sub_column_indices,cols_in_A);
  printf("\nB = A([ ");
  for(int i =0;i<sub_row_indices.size();i++)
    printf("%d ",(int)sub_row_indices[i]);
  printf("],[ ");
  for(int i =0;i<sub_column_indices.size();i++)
    printf("%d ",(int)sub_column_indices[i]);
  printf("]);\n\n");

  printf("B=\n");
  for(int i=0;i<sub_row_indices.size();i++){
    for(int j=0;j<sub_column_indices.size();j++)
      printf("%d ",B[i*sub_column_indices.size()+j]);
    printf("\n");
  }
  
  return 0;
}
Which should output:
A=
10 11 12 
13 14 15 

A[0,2]=12

B = A([ 0 ],[ 2 1 ]);

B=
12 11