Discrete log

Author: adamant-pwn

cp-algo/math/number_theory.hpp

@@ -8,6 +8,67 @@
 #include <vector>
 #include <bit>
 namespace cp_algo::math {
+    std::vector<int64_t> factorize(int64_t m);
+
+    int64_t euler_phi(int64_t m) {
+        auto primes = factorize(m);
+        auto [from, to] = std::ranges::unique(primes);
+        primes.erase(from, to);
+        int64_t ans = m;
+        for(auto it: primes) {
+            ans -= ans / it;
+        }
+        return ans;
+    }
+    template<modint_type base>
+    int64_t period(base x) {
+        auto ans = euler_phi(base::mod());
+        base x0 = bpow(x, ans);
+        for(auto t: factorize(ans)) {
+            while(ans % t == 0 && x0 * bpow(x, ans / t) == x0) {
+                ans /= t;
+            }
+        }
+        return ans;
+    }
+    // Find min non-negative x s.t. a*b^x = c (mod m)
+    std::optional<uint64_t> discrete_log(int64_t b, int64_t c, uint64_t m, int64_t a = 1) {
+        if(std::abs(a - c) % m == 0) {
+            return 0;
+        }
+        if(std::gcd(a, m) != std::gcd(a * b, m)) {
+            auto res = discrete_log(b, c, m, a * b % m);
+            return res ? std::optional(*res + 1) : res;
+        }
+        // a * b^x is periodic here
+        using base = dynamic_modint;
+        return base::with_mod(m, [&]() -> std::optional<uint64_t> {
+            size_t sqrtmod = std::max<size_t>(1, std::sqrt(m) / 2);
+            std::unordered_map<int64_t, int> small;
+            base cur = a;
+            for(size_t i = 0; i < sqrtmod; i++) {
+                small[cur.getr()] = i;
+                cur *= b;
+            }
+            base step = bpow(base(b), sqrtmod);
+            cur = 1;
+            for(size_t k = 0; k < m; k += sqrtmod) {
+                auto it = small.find((base(c) * cur).getr());
+                if(it != end(small)) {
+                    auto cand = base::with_mod(period(base(b)), [&](){
+                        return base(it->second - k);
+                    }).getr();
+                    if(base(a) * bpow(base(b), cand) == base(c)) {
+                        return cand;
+                    } else {
+                        return std::nullopt;
+                    }
+                }
+                cur *= step;
+            }
+            return std::nullopt;
+        });
+    }
     // https://en.wikipedia.org/wiki/Berlekamp-Rabin_algorithm
     template<modint_type base>
     std::optional<base> sqrt(base b) {
@@ -98,17 +159,8 @@ namespace cp_algo::math {
     int64_t primitive_root(int64_t p) {
         using base = dynamic_modint;
         return base::with_mod(p, [p](){
-            auto fact = factorize(p - 1);
-            auto is_primitive_root = [&](base x) {
-                for(auto t: fact) {
-                    if(bpow(x, (p - 1) / t) == 1) {
-                        return false;
-                    }
-                }
-                return true;
-            };
             base t = 1;
-            while(!is_primitive_root(t)) {
+            while(period(t) != p - 1) {
                 t = random::rng();
             }
             return t.getr();

verify/number_theory/discrete_log.test.cpp

@@ -0,0 +1,33 @@
+// @brief Discrete Logarithm
+#define PROBLEM "https://judge.yosupo.jp/problem/discrete_logarithm_mod"
+#pragma GCC optimize("Ofast,unroll-loops")
+#pragma GCC target("avx2,tune=native")
+#include "cp-algo/math/number_theory.hpp"
+#include <bits/stdc++.h>
+
+using namespace std;
+using namespace cp_algo;
+using namespace math;
+using base = dynamic_modint;
+
+void solve() {
+    int x, y, m;
+    cin >> x >> y >> m;
+    auto res = discrete_log(x, y, m);
+    if(res) {
+        cout << *res << "\n";
+    } else {
+        cout << -1 << "\n";
+    }
+}
+
+signed main() {
+    //freopen("input.txt", "r", stdin);
+    ios::sync_with_stdio(0);
+    cin.tie(0);
+    int t = 1;
+    cin >> t;
+    while(t--) {
+        solve();
+    }
+}