feat(math): impl remove zeros for sparse univariate poly

kroma-network · Jun 12, 2024 · 802ef70 · 802ef70
1 parent b3ea3d7
commit 802ef70
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 12 deletions.
diff --git a/tachyon/math/polynomials/univariate/BUILD.bazel b/tachyon/math/polynomials/univariate/BUILD.bazel
@@ -103,6 +103,7 @@ tachyon_cc_library(
         "//tachyon/base/buffer:copyable",
         "//tachyon/base/containers:adapters",
         "//tachyon/base/containers:container_util",
+        "//tachyon/base/containers:cxx20_erase",
         "//tachyon/base/json",
         "//tachyon/base/ranges:algorithm",
         "//tachyon/base/strings:string_util",

diff --git a/tachyon/math/polynomials/univariate/univariate_sparse_coefficients.h b/tachyon/math/polynomials/univariate/univariate_sparse_coefficients.h
@@ -20,6 +20,7 @@
 
 #include "tachyon/base/buffer/copyable.h"
 #include "tachyon/base/containers/adapters.h"
+#include "tachyon/base/containers/cxx20_erase_vector.h"
 #include "tachyon/base/json/json.h"
 #include "tachyon/base/logging.h"
 #include "tachyon/base/optional.h"
@@ -104,13 +105,13 @@ class UnivariateSparseCoefficients {
   constexpr explicit UnivariateSparseCoefficients(
       const std::vector<Term>& terms, bool cleanup = false)
       : terms_(terms) {
-    if (cleanup) RemoveHighDegreeZeros();
+    if (cleanup) RemoveZeros();
     CHECK_LE(Degree(), kMaxDegree);
   }
   constexpr explicit UnivariateSparseCoefficients(std::vector<Term>&& terms,
                                                   bool cleanup = false)
       : terms_(std::move(terms)) {
-    if (cleanup) RemoveHighDegreeZeros();
+    if (cleanup) RemoveZeros();
     CHECK_LE(Degree(), kMaxDegree);
   }
 
@@ -283,6 +284,10 @@ class UnivariateSparseCoefficients {
     return ss.str();
   }
 
+  void RemoveZeros() {
+    base::EraseIf(terms_, [](const Term& x) { return x.coefficient.IsZero(); });
+  }
+
  private:
   friend class internal::UnivariatePolynomialOp<
       UnivariateDenseCoefficients<F, MaxDegree>>;
@@ -299,16 +304,6 @@ class UnivariateSparseCoefficients {
     return &it->coefficient;
   }
 
-  void RemoveHighDegreeZeros() {  // Fix to RemoveZeros
-    while (!IsZero()) {
-      if (terms_.back().coefficient.IsZero()) {
-        terms_.pop_back();
-      } else {
-        break;
-      }
-    }
-  }
-
   std::vector<Term> terms_;
 };