Fixed corner cases

fredbr · fredbr · commit b98f1f72a0b7 · 2019-03-20T16:21:04.000-03:00
diff --git a/math/inverse.cpp b/math/inverse.cpp
@@ -1,4 +1,7 @@
-#include <bits/stdc++.h>
+#include <iostream>
+#include <vector>
+
+#include <gmpxx.h>
 
 template<typename T>
 void equalize_row(std::vector<T> const& A, std::vector<T> &B, T ratio) {
@@ -8,6 +11,13 @@ void equalize_row(std::vector<T> const& A, std::vector<T> &B, T ratio) {
     }
 }
 
+template<typename T>
+void divide_row(std::vector<T> &A, T ratio) {
+    for (T& num : A) {
+        num /= ratio;
+    }
+}
+
 template<typename T>
 void print_matrix(std::vector<T> const& A) {
     for (auto& row : A) {
@@ -27,48 +37,52 @@ std::vector<std::vector<T>> inverse(std::vector<std::vector<T>> M) {
         R[i][i] = T(1);
     }
 
-    for (int32_t i = 0; i+1 < n; i++) {
-        T cut_base = M[i][i];
+    for (int32_t i = 0; i < n; i++) {
+        int32_t best = i;
         for (int32_t j = i+1; j < n; j++) {
-            T ratio = - M[j][i] / cut_base;
-            equalize_row(M[i], M[j], ratio);
-            equalize_row(R[i], R[j], ratio);
+            if (abs(M[best][i]) < abs(M[j][i])) {
+                best = j;
+            }
+        }
+        if (best != i) {
+            std::swap(M[i], M[best]);
+            std::swap(R[i], R[best]);
+        }
+
+        divide_row<T>(R[i], M[i][i]);
+        divide_row<T>(M[i], M[i][i]);
+
+        for (int32_t j = i+1; j < n; j++) {
+            T ratio = -M[j][i];
+            equalize_row<T>(M[i], M[j], ratio);
+            equalize_row<T>(R[i], R[j], ratio);
         }
     }
 
     for (int32_t i = n-1; i > 0; i--) {
-        T cut_base = M[i][i];
         for (int32_t j = i-1; j >= 0; j--) {
-            T ratio = - M[j][i] / cut_base;
-            equalize_row(M[i], M[j], ratio);
-            equalize_row(R[i], R[j], ratio);
+            equalize_row<T>(R[i], R[j], -M[j][i]);
+            equalize_row<T>(M[i], M[j], -M[j][i]);
         }
     }
 
-    print_matrix(M);
-    for (int32_t i = 0; i < n; i++) {
-        if (M[i][i] != T(1)) {
-            T ratio = T(1) / M[i][i] - T(1);
-            equalize_row(R[i], R[i], ratio);
-            equalize_row(M[i], M[i], ratio);
-        }
-    }
-    print_matrix(M);
     return R;
 }
 
 int main() {
+    using rational = mpq_class;
+
     int32_t n;
     std::cin >> n;
 
-    auto M = std::vector<std::vector<double>>(n, std::vector<double>(n));
+    auto M = std::vector<std::vector<rational>>(n, std::vector<rational>(n));
     for (auto& row : M) {
         for (auto& x : row) {
             std::cin >> x;
         }
     }
 
-    auto ans = inverse(M);
+    auto A = inverse<rational>(M);
 
-    print_matrix(ans);
+    print_matrix(A);
 }
