Discrete Logarithm (verify/number_theory/discrete_log.test.cpp)

View this file on GitHub
Last update: 2024-10-07 21:21:13+02:00
Problem: https://judge.yosupo.jp/problem/discrete_logarithm_mod

Depends on

Code

// @brief Discrete Logarithm
#define PROBLEM "https://judge.yosupo.jp/problem/discrete_logarithm_mod"
#pragma GCC optimize("Ofast,unroll-loops")
#include "cp-algo/number_theory/discrete_log.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();
    }
}

#line 1 "verify/number_theory/discrete_log.test.cpp"
// @brief Discrete Logarithm
#define PROBLEM "https://judge.yosupo.jp/problem/discrete_logarithm_mod"
#pragma GCC optimize("Ofast,unroll-loops")
#line 1 "cp-algo/number_theory/discrete_log.hpp"


#line 1 "cp-algo/number_theory/euler.hpp"


#line 1 "cp-algo/number_theory/factorize.hpp"


#line 1 "cp-algo/number_theory/primality.hpp"


#line 1 "cp-algo/number_theory/modint.hpp"


#line 1 "cp-algo/math/common.hpp"


#include <functional>
#include <cstdint>
namespace cp_algo::math {
#ifdef CP_ALGO_MAXN
    const int maxn = CP_ALGO_MAXN;
#else
    const int maxn = 1 << 19;
#endif
    const int magic = 64; // threshold for sizes to run the naive algo

    auto bpow(auto const& x, int64_t n, auto const& one, auto op) {
        if(n == 0) {
            return one;
        } else {
            auto t = bpow(x, n / 2, one, op);
            t = op(t, t);
            if(n % 2) {
                t = op(t, x);
            }
            return t;
        }
    }
    auto bpow(auto x, int64_t n, auto ans) {
        return bpow(x, n, ans, std::multiplies{});
    }
    template<typename T>
    T bpow(T const& x, int64_t n) {
        return bpow(x, n, T(1));
    }
}

#line 4 "cp-algo/number_theory/modint.hpp"
#include <iostream>
#include <cassert>
namespace cp_algo::math {
    inline constexpr uint64_t inv64(uint64_t x) {
        assert(x % 2);
        uint64_t y = 1;
        while(y * x != 1) {
            y *= 2 - x * y;
        }
        return y;
    }

    template<typename modint>
    struct modint_base {
        static int64_t mod() {
            return modint::mod();
        }
        static uint64_t imod() {
            return modint::imod();
        }
        static __uint128_t pw128() {
            return modint::pw128();
        }
        static uint64_t m_reduce(__uint128_t ab) {
            if(mod() % 2 == 0) [[unlikely]] {
                return ab % mod();
            } else {
                uint64_t m = ab * imod();
                return (ab + __uint128_t(m) * mod()) >> 64;
            }
        }
        static uint64_t m_transform(uint64_t a) {
            if(mod() % 2 == 0) [[unlikely]] {
                return a;
            } else {
                return m_reduce(a * pw128());
            }
        }
        modint_base(): r(0) {}
        modint_base(int64_t rr): r(rr % mod()) {
            r = std::min(r, r + mod());
            r = m_transform(r);
        }
        modint inv() const {
            return bpow(to_modint(), mod() - 2);
        }
        modint operator - () const {
            modint neg;
            neg.r = std::min(-r, 2 * mod() - r);
            return neg;
        }
        modint& operator /= (const modint &t) {
            return to_modint() *= t.inv();
        }
        modint& operator *= (const modint &t) {
            r = m_reduce(__uint128_t(r) * t.r);
            return to_modint();
        }
        modint& operator += (const modint &t) {
            r += t.r; r = std::min(r, r - 2 * mod());
            return to_modint();
        }
        modint& operator -= (const modint &t) {
            r -= t.r; r = std::min(r, r + 2 * mod());
            return to_modint();
        }
        modint operator + (const modint &t) const {return modint(to_modint()) += t;}
        modint operator - (const modint &t) const {return modint(to_modint()) -= t;}
        modint operator * (const modint &t) const {return modint(to_modint()) *= t;}
        modint operator / (const modint &t) const {return modint(to_modint()) /= t;}
        // Why <=> doesn't work?..
        auto operator == (const modint_base &t) const {return getr() == t.getr();}
        auto operator != (const modint_base &t) const {return getr() != t.getr();}
        auto operator <= (const modint_base &t) const {return getr() <= t.getr();}
        auto operator >= (const modint_base &t) const {return getr() >= t.getr();}
        auto operator < (const modint_base &t) const {return getr() < t.getr();}
        auto operator > (const modint_base &t) const {return getr() > t.getr();}
        int64_t rem() const {
            uint64_t R = getr();
            return 2 * R > (uint64_t)mod() ? R - mod() : R;
        }

        // Only use if you really know what you're doing!
        uint64_t modmod() const {return 8ULL * mod() * mod();};
        void add_unsafe(uint64_t t) {r += t;}
        void pseudonormalize() {r = std::min(r, r - modmod());}
        modint const& normalize() {
            if(r >= (uint64_t)mod()) {
                r %= mod();
            }
            return to_modint();
        }
        void setr(uint64_t rr) {r = m_transform(rr);}
        uint64_t getr() const {
            uint64_t res = m_reduce(r);
            return std::min(res, res - mod());
        }
        void setr_direct(uint64_t rr) {r = rr;}
        uint64_t getr_direct() const {return std::min(r, r - mod());}
    private:
        uint64_t r;
        modint& to_modint() {return static_cast<modint&>(*this);}
        modint const& to_modint() const {return static_cast<modint const&>(*this);}
    };
    template<typename modint>
    std::istream& operator >> (std::istream &in, modint_base<modint> &x) {
        uint64_t r;
        auto &res = in >> r;
        x.setr(r);
        return res;
    }
    template<typename modint>
    std::ostream& operator << (std::ostream &out, modint_base<modint> const& x) {
        return out << x.getr();
    }

    template<typename modint>
    concept modint_type = std::is_base_of_v<modint_base<modint>, modint>;

    template<int64_t m>
    struct modint: modint_base<modint<m>> {
        static constexpr uint64_t im = m % 2 ? inv64(-m) : 0;
        static constexpr uint64_t r2 = __uint128_t(-1) % m + 1;
        static constexpr int64_t mod() {return m;}
        static constexpr uint64_t imod() {return im;}
        static constexpr __uint128_t pw128() {return r2;}
        using Base = modint_base<modint<m>>;
        using Base::Base;
    };

    struct dynamic_modint: modint_base<dynamic_modint> {
        static int64_t mod() {return m;}
        static uint64_t imod() {return im;}
        static __uint128_t pw128() {return r2;}
        static void switch_mod(int64_t nm) {
            m = nm;
            im = m % 2 ? inv64(-m) : 0;
            r2 = __uint128_t(-1) % m + 1;
        }
        using Base = modint_base<dynamic_modint>;
        using Base::Base;

        // Wrapper for temp switching
        auto static with_mod(int64_t tmp, auto callback) {
            struct scoped {
                int64_t prev = mod();
                ~scoped() {switch_mod(prev);}
            } _;
            switch_mod(tmp);
            return callback();
        }
    private:
        static int64_t m;
        static uint64_t im, r1, r2;
    };
    int64_t dynamic_modint::m = 1;
    uint64_t dynamic_modint::im = -1;
    uint64_t dynamic_modint::r2 = 0;
}

#line 4 "cp-algo/number_theory/primality.hpp"
#include <algorithm>
namespace cp_algo::math {
    // https://en.wikipedia.org/wiki/Miller–Rabin_primality_test
    bool is_prime(uint64_t m) {
        if(m == 1 || m % 2 == 0) {
            return m == 2;
        }
        // m - 1 = 2^s * d
        int s = std::countr_zero(m - 1);
        auto d = (m - 1) >> s;
        using base = dynamic_modint;
        auto test = [&](base x) {
            x = bpow(x, d);
            if(std::abs(x.rem()) <= 1) {
                return true;
            }
            for(int i = 1; i < s && x != -1; i++) {
                x *= x;
            }
            return x == -1;
        };
        return base::with_mod(m, [&](){
            // Works for all m < 2^64: https://miller-rabin.appspot.com
            return std::ranges::all_of(std::array{
                2, 325, 9375, 28178, 450775, 9780504, 1795265022
            }, test);
        });
    }
}

#line 1 "cp-algo/random/rng.hpp"


#include <chrono>
#include <random>
namespace cp_algo::random {
    uint64_t rng() {
        static std::mt19937_64 rng(
            std::chrono::steady_clock::now().time_since_epoch().count()
        );
        return rng();
    }
}

#line 5 "cp-algo/number_theory/factorize.hpp"
namespace cp_algo::math {
    // https://en.wikipedia.org/wiki/Pollard%27s_rho_algorithm
    void factorize(uint64_t m, std::vector<int64_t> &res) {
        if(m % 2 == 0) {
            factorize(m / 2, res);
            res.push_back(2);
        } else if(is_prime(m)) {
            res.push_back(m);
        } else if(m > 1) {
            using base = dynamic_modint;
            base::with_mod(m, [&]() {
                base t = random::rng();
                auto f = [&](auto x) {
                    return x * x + t;
                };
                base x, y;
                base g = 1;
                while(g == 1) {
                    for(int i = 0; i < 64; i++) {
                        x = f(x);
                        y = f(f(y));
                        if(x == y) [[unlikely]] {
                            t = random::rng();
                            x = y = 0;
                        } else {
                            base t = g * (x - y);
                            g = t == 0 ? g : t;
                        }
                    }
                    g = std::gcd(g.getr(), m);
                }
                factorize(g.getr(), res);
                factorize(m / g.getr(), res);
            });
        }
    }
    std::vector<int64_t> factorize(int64_t m) {
        std::vector<int64_t> res;
        factorize(m, res);
        return res;
    }
}

#line 4 "cp-algo/number_theory/euler.hpp"
namespace cp_algo::math {
    int64_t euler_phi(int64_t m) {
        auto primes = factorize(m);
        std::ranges::sort(primes);
        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;
    }
    int64_t primitive_root(int64_t p) {
        using base = dynamic_modint;
        return base::with_mod(p, [p](){
            base t = 1;
            while(period(t) != p - 1) {
                t = random::rng();
            }
            return t.getr();
        });
    }
}

#line 4 "cp-algo/number_theory/discrete_log.hpp"
namespace cp_algo::math {
    // 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;
        });
    }
}

#line 5 "verify/number_theory/discrete_log.test.cpp"
#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();
    }
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	even_mod_00	AC	300 ms	4 MB
g++	even_mod_01	AC	298 ms	4 MB
g++	even_mod_impossible_00	AC	282 ms	4 MB
g++	even_mod_impossible_01	AC	283 ms	4 MB
g++	example_00	AC	5 ms	4 MB
g++	max_random_00	AC	277 ms	4 MB
g++	max_random_01	AC	274 ms	4 MB
g++	max_random_02	AC	265 ms	4 MB
g++	max_random_yes_00	AC	134 ms	4 MB
g++	max_random_yes_01	AC	134 ms	4 MB
g++	max_random_yes_prime_00	AC	174 ms	4 MB
g++	max_random_yes_prime_01	AC	189 ms	4 MB
g++	random_00	AC	96 ms	4 MB
g++	random_01	AC	133 ms	4 MB
g++	random_02	AC	163 ms	4 MB
g++	random_prime_00	AC	227 ms	4 MB
g++	random_prime_01	AC	222 ms	4 MB
g++	small_00	AC	5 ms	4 MB
g++	small_01	AC	5 ms	4 MB
g++	small_02	AC	5 ms	4 MB