ENH: More speedups on Hungarian

WeatherGod · May 17, 2011 · 562d5b3 · 562d5b3
1 parent f805703
commit 562d5b3
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/scikits/learn/utils/hungarian.py b/scikits/learn/utils/hungarian.py
@@ -126,11 +126,11 @@ def _step4(self):
         """
         # We convert to int as numpy operations are faster on int
         C = (self.C == 0).astype(np.int)
+        covered_C = C*self.row_uncovered[:, np.newaxis]
+        covered_C *= self.col_uncovered.astype(np.int)
         n = self.n
         while True:
             # Find an uncovered zero
-            covered_C = C*self.row_uncovered[:, np.newaxis]
-            covered_C *= self.col_uncovered.astype(np.int)
             raveled_idx = np.argmax(covered_C)
             col = raveled_idx % n
             row = raveled_idx // n
@@ -149,6 +149,10 @@ def _step4(self):
                     col = star_col
                     self.row_uncovered[row] = False
                     self.col_uncovered[col] = True
+                    covered_C[:, col] = C[:, col]*(
+                                self.row_uncovered.astype(np.int))
+                    covered_C[row] = 0
+
 
     def _step5(self):
         """