Sum of Totien Function

Author: adamant-pwn

cp-algo/number_theory/dirichlet.hpp

@@ -10,6 +10,18 @@ namespace cp_algo::math {
         return std::pair{rt_n, 2 * rt_n - (n / rt_n == rt_n)};
     }
 
+    auto floor_generator(int64_t n) {
+        auto [rt_n, num_floors] = floor_stats(n);
+        return [n, rt_n = rt_n, num_floors = num_floors](int k) {
+            return k <= rt_n ? int64_t(k) : n / int64_t(num_floors - k + 1);
+        };
+    }
+
+    auto floors(int64_t n) {
+        auto [_, m] = floor_stats(n);
+        return std::views::iota(0, m+1) | std::views::transform(floor_generator(n));
+    }
+
     struct interval {
         int lo, hi;
         auto operator <=>(const interval&) const = default;
@@ -96,8 +108,8 @@ namespace cp_algo::math {
     }
 
     auto Dirichlet_div(auto const& H, auto const& G, int64_t n) {
-        auto m = size(G);
-        auto F = H | std::views::take(m) | std::ranges::to<std::decay_t<decltype(G)>>();
+        auto m = std::size(G);
+        auto F = H | std::views::take(m) | std::ranges::to<std::vector>();
         Dirichlet_div_inplace(F, G, n);
         return F;
     }

verify/number_theory/totient_sum.test.cpp

@@ -0,0 +1,37 @@
+// @brief Sum of Totient Function
+#define PROBLEM "https://judge.yosupo.jp/problem/sum_of_totient_function"
+#pragma GCC optimize("Ofast,unroll-loops")
+#include <iostream>
+#include "blazingio/blazingio.min.hpp"
+#include "cp-algo/number_theory/modint.hpp"
+#include "cp-algo/number_theory/dirichlet.hpp"
+#include "cp-algo/util/big_alloc.hpp"
+#include <bits/stdc++.h>
+
+using namespace std;
+using namespace cp_algo::math;
+using base = modint<998244353>;
+
+void solve() {
+    int64_t n;
+    cin >> n;
+    auto H = floors(n) | views::transform([](int64_t k) {
+        return base(k) * base(k + 1) / 2;
+    }) | ranges::to<vector>();
+    auto G = floors(n) | views::transform([](int64_t k) {
+        return base(k);
+    }) | ranges::to<vector>();
+    auto F = Dirichlet_div(H, G, n);
+    cout << F.back() << "\n";
+}
+
+signed main() {
+    //freopen("input.txt", "r", stdin);
+    ios::sync_with_stdio(0);
+    cin.tie(0);
+    int t = 1;
+    //cin >> t;
+    while(t--) {
+        solve();
+    }
+}