refactored

2 files changed, 15 insertions, 18 deletions

diff --git a/min_plus_conv.h b/min_plus_conv.h
index 2b5e192..b024e7f 100644
--- a/min_plus_conv.h
+++ b/min_plus_conv.h
@@ -3,12 +3,12 @@
 #include <limits>
 #include <vector>
 
-template <typename E_> struct MinPlusConv {
+template <typename T> struct MinPlusConv {
   struct Monge {
-    using E = E_;
+    using E = T;
 
-    E operator()(int k, int i) const {
-      constexpr E inf = std::numeric_limits<E>::max();
+    T operator()(int k, int i) const {
+      constexpr T inf = std::numeric_limits<T>::max();
       auto j = k - i;
       if (0 <= j && j < w.size()) {
         return a[i] + w[k - i];
@@ -16,20 +16,14 @@ template <typename E_> struct MinPlusConv {
       return j < 0 ? inf : inf - i;
     }
 
-    const std::vector<E> &a, &w;
+    const std::vector<T> &a, &w;
   };
 
-  using Seq = std::vector<E_>;
+  using Seq = std::vector<T>;
 
   // Assume that w[i] - w[i - 1] <= w[i + 1] - w[i]
   Seq operator()(const Seq &a, const Seq &w) {
-    int n = a.size() + w.size() - 1;
-    auto row_min = smawk(Monge{a, w}, n, a.size());
-    Seq c(n);
-    for (int i = 0; i < n; i++) {
-      c[i] = row_min[i].first;
-    }
-    return c;
+    return smawk(Monge{a, w}, a.size() + w.size() - 1, a.size());
   }
 
   SMAWK<Monge> smawk;
diff --git a/smawk.h b/smawk.h
index c22b332..0fe3018 100644
--- a/smawk.h
+++ b/smawk.h
@@ -6,13 +6,14 @@
 template <IsTM A> struct SMAWK {
   using E = typename A::E;
   using EI = std::pair<E, int>;
-  using Result = std::vector<EI>;
+  using Result = std::vector<E>;
 
   void operator()(Result &row_min, const A &a, int n, int m) {
     stack.resize(n);
     cols.resize(m + n + n);
     std::iota(cols.begin(), cols.begin() + m, 0);
     row_min.resize(n);
+    row_argmin.resize(n);
     recur(row_min, a, n, m, 0, 0, m);
   }
 
@@ -26,10 +27,11 @@ template <IsTM A> struct SMAWK {
              int end) {
     if (n < (2 << k)) {
       auto r = (1 << k) - 1;
-      auto &ref = row_min[r] = query(a, r, cols[begin]);
+      auto ref = query(a, r, cols[begin]);
       for (int i = begin + 1; i < end; i++) {
         ref = std::min(ref, query(a, r, cols[i]));
       }
+      std::tie(row_min[r], row_argmin[r]) = ref;
     } else {
       int top{0};
       for (int i = begin; i < end; i++) {
@@ -47,11 +49,12 @@ template <IsTM A> struct SMAWK {
       recur(row_min, a, n, m, k + 1, begin, end);
       auto offset = 1 << k;
       for (int r = offset - 1, p = begin; r < n; r += offset << 1) {
-        auto high = r + offset < n ? row_min[r + offset].second + 1 : m;
-        auto &ref = row_min[r] = query(a, r, cols[p]);
+        auto high = r + offset < n ? row_argmin[r + offset] + 1 : m;
+        auto ref = query(a, r, cols[p]);
         while (p + 1 < end && cols[p + 1] < high) {
           ref = std::min(ref, query(a, r, cols[++p]));
         }
+        std::tie(row_min[r], row_argmin[r]) = ref;
       }
     }
   }
@@ -59,5 +62,5 @@ template <IsTM A> struct SMAWK {
   EI query(const A &a, int r, int c) { return {a(r, c), c}; }
 
   std::vector<E> stack;
-  std::vector<int> cols;
+  std::vector<int> cols, row_argmin;
 };