Simplify Dirichlet a bit

Author: adamant-pwn

cp-algo/number_theory/dirichlet.hpp

@@ -15,9 +15,9 @@ namespace cp_algo::math {
         auto operator <=>(const interval&) const = default;
     };
 
-    // callback(k, prefix) such that:
-    //     (F * G)[k] = prefix + (F[k] - F[k-1]) * G[1] + (G[k] - G[k-1]) * F[1]
-    // Uses H as a buffer for (F * G)[k], then overrides with callback results
+    // callback(k) when:
+    //     (F * G)[k] = H[k] + (F[k] - F[k-1]) * G[1] + (G[k] - G[k-1]) * F[1]
+    // Return the value to be saved in H[k]
     enum exec_mode { standard, reverse };
     template<exec_mode mode = standard>
     void exec_on_blocks(int64_t n, auto &H, auto const& F, auto const& G, auto &&callback) {
@@ -30,26 +30,23 @@ namespace cp_algo::math {
         auto call = [&](interval x, interval y, interval z) {
             auto Fx = F[x.hi] - F[x.lo - 1];
             auto Fy = F[y.hi] - F[y.lo - 1];
-            decltype(Fx) Gx, Gy, t;
+            decltype(Fx) Gx, Gy;
             if constexpr (mode == standard) {
                 Gy = G[y.hi] - G[y.lo - 1];
                 Gx = G[x.hi] - G[x.lo - 1];
             } else {
                 Gy = G[y.lo - 1] - G[y.hi];
                 Gx = G[x.lo - 1] - G[x.hi];
             }
-            if(x == y) [[unlikely]] {
-                t = Fx * Gy;
-            } else {
-                t = Fx * Gy + Fy * Gx;
+            auto t = Fx * Gy;
+            if(x != y) [[likely]] {
+                t += Fy * Gx;
             }
             H[z.lo] += t;
-            if (z.hi < num_floors)  {
+            if (z.hi < num_floors) [[likely]] {
                 H[z.hi + 1] -= t;
             }
         };
-
-        auto prefix = F[1] * G[1];
         for (int k = 2; k <= num_floors; ++k) {
             if(k > rt_n) {
                 int z = num_floors - k + 1;
@@ -61,8 +58,7 @@ namespace cp_algo::math {
                 }
             }
 
-            H[k] = callback(k, prefix += H[k]);
-            prefix += (F[k] - F[k-1]) * G[1] + (G[k] - G[k-1]) * F[1];
+            H[k] = callback(k);
 
             if(k <= rt_n) {
                 int x = k;
@@ -82,9 +78,10 @@ namespace cp_algo::math {
         auto m = size(F);
         std::decay_t<decltype(F)> H(m);
         H[1] = F[1] * G[1];
-        exec_on_blocks(n, H, F, G, [&](auto k, auto prefix) {
-            return prefix + (F[k] - F[k-1]) * G[1] + (G[k] - G[k-1]) * F[1];
+        exec_on_blocks(n, H, F, G, [&](auto k) {
+            return H[k] + (F[k] - F[k-1]) * G[1] + (G[k] - G[k-1]) * F[1];
         });
+        partial_sum(begin(H), end(H), begin(H));
         return H;
     }
 
@@ -94,7 +91,7 @@ namespace cp_algo::math {
         H[0] -= H[0];
         adjacent_difference(begin(H), end(H), begin(H));
         H[1] *= Gi;
-        exec_on_blocks<reverse>(n, H, H, G, [&](auto k, auto) {
+        exec_on_blocks<reverse>(n, H, H, G, [&](auto k) {
             return (Gi * (H[k] - (G[k] - G[k-1]) * H[1])) + H[k-1];
         });
     }