src/hungarian.cpp

   1 #include <iostream>
   2 #include <algorithm>
   3 #include "hungarian.h"
   4
   5 using namespace std;
   6 using namespace Eigen;
   7
   8 class Hungary;
   9
  10 int step_one(Hungary& state);
  11 int step_two(Hungary& state);
  12 int step_three(Hungary& state);
  13 int step_four(Hungary& state);
  14 int step_five(Hungary& state);
  15 int step_six(Hungary& state);
  16
  17 class Hungary
  18 {
  19 public:
  20     Hungary(const MatrixXi& cost){
  21         this->cost = cost;
  22         this->row_uncovered = VectorXi::Ones(cost.rows());
  23         this->col_uncovered = VectorXi::Ones(cost.cols());
  24         Z0_r = 0;
  25         Z0_c = 0;
  26         this->mark = MatrixXi::Zero(cost.rows(), cost.cols());
  27     };
  28     ~Hungary(){};
  29
  30     void print(){
  31         cout <<"Cost:\n" << cost << "\nMark:\n" << mark << endl;
  32         cout <<"Row uncovered :[" << row_uncovered.transpose() << "], col uncovered: [" << col_uncovered.transpose() << "]" << endl;
  33     }
  34
  35 public:
  36     void clearCovers(){
  37         this->row_uncovered.setOnes();
  38         this->col_uncovered.setOnes();
  39     }
  40
  41 public :
  42     // The cost matrix, cols() should > rows(). If not, transpose it first.
  43     MatrixXi cost;
  44     // The mark matrix. the value of the element is 0 (None), 1 (means the starred), 2 (primed)
  45     MatrixXi mark;
  46     // the covered state of the rows (0: covered; 1: nocovered)
  47     VectorXi row_uncovered;
  48     // the covered state of the column (0: covered; 1: nocovered)
  49     VectorXi col_uncovered;
  50     // the position of Z0, used in step 5
  51     int Z0_r;
  52     int Z0_c;
  53     int min_value;
  54 };
  55
  56 int linear_sum_assignment(const MatrixXi& cost_matrix, VectorXi& row_ind, VectorXi& col_ind)
  57 {
  58     // The algorithm expects more columns than rows in the cost matrix.
  59     MatrixXi correct_matrix = cost_matrix;
  60     if (cost_matrix.cols() == 0 || cost_matrix.rows() == 0){
  61         row_ind = VectorXi::Zero(0);
  62         col_ind = VectorXi::Zero(0);
  63         return 0;
  64     }
  65     bool is_transposed = false;
  66     if (cost_matrix.cols() < cost_matrix.rows()){
  67         cout << "cols < rows, transpose." << endl;
  68         correct_matrix = cost_matrix.transpose();
  69         is_transposed = true;
  70     }
  71     Hungary state(correct_matrix);
  72     cout << "Cost Matrix: \n" << correct_matrix << endl;;
  73
  74     bool done = false;
  75     int next = 1;
  76     int pre_step = 1;
  77     while (!done){
  78         pre_step = next;
  79         switch(next)
  80         {
  81             case 1:
  82                 next = step_one(state);
  83                 break;
  84             case 2:
  85                 next = step_two(state);
  86                 break;
  87             case 3:
  88                 next = step_three(state);
  89                 break;
  90             case 4:
  91                 next = step_four(state);
  92                 break;
  93             case 5:
  94                 next = step_five(state);
  95                 break;
  96             case 6:
  97                 next = step_six(state);
  98                 break;
  99             case 7:
 100                 done = true;
 101                 break;
 102         }
 103         cout << "After step: " << pre_step << endl;
 104         state.print();
 105     }
 106     MatrixXi mark = state.mark;
 107     if (is_transposed){
 108         mark = state.mark.transpose();
 109     }
 110     cout << "Done" << endl << mark << endl;
 111     int sum = (mark.array() * cost_matrix.array()).sum();
 112     cout << "Sum :" << sum << endl;
 113
 114
 115     // return the array of the positions
 116     row_ind = VectorXi::Zero(mark.cols());
 117     col_ind = VectorXi::Zero(mark.cols());
 118     VectorXi::Index max_index;
 119     int point = 0;
 120     for (int i = 0; i < mark.rows(); i++){
 121         if (mark.row(i).maxCoeff(&max_index)){
 122             row_ind(point) = i;
 123             col_ind(point) = max_index;
 124             point++;
 125         }
 126     }
 127     return sum;
 128 }
 129
 130
 131 // Step 1:
 132 // For each row of the matrix, find the smallest element and subtract
 133 // it from every element in its row. Go to Step 2.
 134 int step_one(Hungary& state)
 135 {
 136     state.cost.colwise() -= state.cost.rowwise().minCoeff();
 137     return 2;
 138 }
 139
 140 // Step 2:
 141 // Find a zero (Z) in the resulting matrix. If there is no starred zero in its row or column,
 142 // star Z. Repeat for each elements in the matrix. Go to Step 3.
 143 int step_two(Hungary& state)
 144 {
 145     for (int i = 0; i < state.cost.rows(); i++){
 146         for (int j = 0; j < state.cost.cols(); j++)
 147             if (state.row_uncovered(i) == 1 && state.col_uncovered(j) == 1 && state.cost(i, j) == 0){
 148                 state.mark(i, j) = 1;
 149                 state.col_uncovered(j) = 0;
 150                 state.row_uncovered(i) = 0;
 151             }
 152     }
 153     state.clearCovers();
 154     return 3;
 155 }
 156
 157 // Step 3:
 158 // Cover each column containing a starred zero. If K columns are covered,
 159 // the starred zeros describe a complete set of unique assignments. In this case,
 160 // Go to DONE, otherwise, Go to Step 4.
 161 int step_three(Hungary& state)
 162 {
 163     MatrixXi m1 = (state.mark.array() == 1).select(state.mark, 0);
 164     state.col_uncovered = m1.colwise().any();
 165     //for (int i = 0; i < state.col_uncovered.size(); i++)
 166     //    if (m1.col(i).any())
 167     //        state.col_uncovered(i) = 0;
 168         //state.col_uncovered(i) = (state.col_uncovered(i) == 1) ? 0 : 1;
 169     state.col_uncovered = (state.col_uncovered.array() == 1).select(0, VectorXi::Ones(state.col_uncovered.size()));
 170     int next = 4;
 171     if (m1.sum() == m1.rows())
 172         next = 7;
 173     return next;
 174 }
 175
 176 // Step 4 :
 177 // Find a noncovered zero and prime it. If there is no starred zero in the row
 178 // constaining this primed zero, Go to Step 5. Otherwise, cover this row and uncover the column
 179 // containing the starred zero. Continue in this manner until there are no uncovered zeros left.
 180 // Save the smallest uncovered value and Go to Step 6.
 181 //
 182 //
 183 // find_uncovered_zero
 184 //   - find zero or minimal value in the uncovered elements.
 185 //   Return 0 if zero found, or the minimal element.
 186 int find_uncovered_zero(Hungary& state, int& row, int& col)
 187 {
 188     MatrixXi cover = state.row_uncovered * state.col_uncovered.transpose();
 189     int min_value = 0;
 190     for (int i = 0; i < cover.rows(); i++)
 191         for (int j = 0; j < cover.cols(); j++){
 192             if (cover(i, j) != 0){
 193                 if (state.cost(i, j) == 0){
 194                     row = i;
 195                     col = j;
 196                     return 0;
 197                 } else {
 198                     min_value = (min_value > 0) ? min(min_value, state.cost(i, j)) : state.cost(i, j);
 199                 }
 200             }
 201         }
 202     return min_value;
 203 }
 204
 205 int step_four(Hungary& state)
 206 {
 207     while(true)
 208     {
 209         int rr = 0;
 210         int cc = 0;
 211         // find a nocovered zero
 212         int min = find_uncovered_zero(state, rr, cc);
 213         if(min == 0){
 214             // prime it
 215             state.mark(rr, cc) = 2;
 216             // If no starred zero in its row
 217             if ((state.mark.row(rr).array() == 1).any() == 0){
 218                 state.Z0_r = rr;
 219                 state.Z0_c = cc;
 220                 return 5;
 221             } else {
 222                 // Otherwise, cover this row
 223                 state.row_uncovered(rr) = 0;
 224                 // uncover the column
 225                 for (int j = 0; j < state.mark.cols(); j++){
 226                     if (state.mark(rr, j) == 1){
 227                         state.col_uncovered(j) = 1;
 228                     }
 229                 }
 230             }
 231         } else {
 232             state.min_value = min;
 233             return 6;
 234         }
 235     }
 236 }
 237
 238 // Step 5:
 239 //   Construct a seriee of alternating primed and starred zeros as follows,
 240 //   Let Z0 represent the uncovered primed zero found in Step 4.
 241 //   Let Z1 denote the starred zero in the column of Z0 (if any).
 242 //   Let Z2 denote the primed zero in the row of Z1 (there will always be one).
 243 //   Continue until the series terminates at a primed zero that has no starred
 244 //   zero in its column. Unstar each starred zero of the series, star each primed
 245 //   zero of the series, erase all primes and uncover every line in the matrix. Return to Step 3
 246 int step_five(Hungary& state)
 247 {
 248     MatrixXi path = MatrixXi::Zero(state.cost.rows() + state.cost.cols(), 2);
 249     int count = 0;
 250     path(count, 0) = state.Z0_r;
 251     path(count, 1) = state.Z0_c;
 252
 253     MatrixXi m1 = (state.mark.array() == 1).select(state.mark, 0);
 254     MatrixXi m2 = (state.mark.array() == 2).select(state.mark, 0);
 255     MatrixXi::Index index;
 256     while(true){
 257         // Z1, find the starred zero in the column of Z0
 258         int max = m1.col(path(count, 1)).maxCoeff(&index);
 259         if (max != 1)
 260             break;
 261         else {
 262             count++;
 263             path(count, 0) = index;
 264             path(count, 1) = path(count - 1, 1);
 265         }
 266
 267         // Z2, find the primed zero in the row of Z1;
 268         max = m2.row(path(count, 0)).maxCoeff(&index);
 269         if (max != 2){
 270             cout << "Error, should never reach here" << endl;
 271         } else {
 272             count++;
 273             path(count, 0) = path(count - 1, 0);
 274             path(count, 1) = index;
 275         }
 276     }
 277
 278     // terminates
 279     for (int i = 0; i < count + 1; i++){
 280         if (state.mark(path(i, 0), path(i, 1)) == 1)
 281             state.mark(path(i, 0), path(i, 1)) = 0;
 282         else
 283             state.mark(path(i, 0), path(i, 1)) = 1;
 284     }
 285
 286     state.clearCovers();
 287
 288     // erase all prime markings
 289     state.mark = (state.mark.array() == 2).select(0, state.mark);
 290     return 3;
 291 }
 292
 293 // Step 6:
 294 //    Add the value found in Step 4 to every element of each covered row, and substract it
 295 //    from every element of each uncovered column.
 296 //    Return to Step 4 without altering any stars, primes, or covered lines.
 297 //
 298 int step_six(Hungary& state)
 299 {
 300     for (int i = 0; i < state.mark.rows(); i++)
 301         for (int j = 0; j < state.mark.cols(); j++){
 302             if (state.row_uncovered(i) == 0)
 303                 state.cost(i, j) += state.min_value;
 304             if (state.col_uncovered(j) == 1)
 305                 state.cost(i, j) -= state.min_value;
 306         }
 307     return 4;
 308 }
 309