Update discrete logarithm

number/discrete_logarithm.hpp

@@ -1,62 +1,92 @@
 #pragma once
-#include <algorithm>
-#include <unordered_map>
-#include <utility>
-
-// CUT begin
-// Calculate log_A B (MOD M) (baby-step gian-step)
-// DiscreteLogarithm dl(M, A);
-// lint ans = dl.log(B);
-// Complexity: O(M^(1/2)) for each query
-// Verified: <https://judge.yosupo.jp/problem/discrete_logarithm_mod>
-// Constraints: 0 <= A < M, B < M, 1 <= M <= 1e9  (M is not limited to prime)
-struct DiscreteLogarithm {
-    using lint = long long int;
-    int M, stepsize;
-    lint baby_a, giant_a, g;
-    std::unordered_map<lint, int> baby_log_dict;
-
-    lint inverse(lint a) {
-        lint b = M / g, u = 1, v = 0;
-        while (b) {
-            lint t = a / b;
-            a -= t * b;
-            std::swap(a, b);
-            u -= t * v;
-            std::swap(u, v);
+#include <cassert>
+#include <cmath>
+#include <functional>
+#include <unordered_set>
+
+// Solve min_n f^n s = t (0 <= n <= max_search)
+// baby = f, giant = f^giant_stride
+// If solution is not found, return -1.
+// https://maspypy.com/%e3%83%a2%e3%83%8e%e3%82%a4%e3%83%89%e4%bd%9c%e7%94%a8%e3%81%ab%e9%96%a2%e3%81%99%e3%82%8b%e9%9b%a2%e6%95%a3%e5%af%be%e6%95%b0%e5%95%8f%e9%a1%8c
+template <class S, class F, class Container>
+long long
+DiscreteLogarithm(const F &baby, const F &giant, const S &s, const S &t,
+                  const std::function<S(F, S)> &mapping, long long max_search, int giant_stride) {
+    if (s == t) return 0;
+    if (max_search <= 0) return -1;
+
+    Container ys;
+    // ys.reserve(giant_stride); // if unordered_set
+    {
+        auto yt = t;
+        for (int i = 0; i < giant_stride; ++i) {
+            ys.emplace(yt);
+            yt = mapping(baby, yt);
         }
-        u %= M / g;
-        return u >= 0 ? u : u + M / g;
     }
 
-    DiscreteLogarithm(int mod, int a_new) : M(mod), baby_a(a_new % mod), giant_a(1) {
-        g = 1;
-        while (std::__gcd(baby_a, M / g) > 1) g *= std::__gcd(baby_a, M / g);
-        stepsize = 32; // lg(MAX_M)
-        while (stepsize * stepsize < M / g) stepsize++;
-
-        lint now = 1 % (M / g), inv_g = inverse(baby_a);
-        for (int n = 0; n < stepsize; n++) {
-            if (!baby_log_dict.count(now)) baby_log_dict[now] = n;
-            (now *= baby_a) %= M / g;
-            (giant_a *= inv_g) %= M / g;
+    int num_fails = 0;
+    S cur = s;
+
+    for (long long k = 1;; ++k) {
+        if (const S nxt = mapping(giant, cur); ys.count(nxt)) {
+            for (int i = 1; i <= giant_stride; ++i) {
+                cur = mapping(baby, cur);
+                if (cur == t) {
+                    long long ret = (k - 1) * giant_stride + i;
+                    return (ret <= max_search) ? ret : -1;
+                }
+            }
+            ++num_fails;
+        } else {
+            cur = nxt;
         }
+
+        if (num_fails >= 2 or k * giant_stride > max_search) return -1;
     }
+}
 
-    // log(): returns the smallest nonnegative x that satisfies a^x = b mod M, or -1 if there's no solution
-    lint log(lint b) {
-        b %= M;
-        lint acc = 1 % M;
-        for (int i = 0; i < stepsize; i++) {
-            if (acc == b) return i;
-            (acc *= baby_a) %= M;
-        }
-        if (b % g) return -1; // No solution
-        lint now = b * giant_a % (M / g);
-        for (lint q = 1; q <= M / stepsize + 1; q++) {
-            if (baby_log_dict.count(now)) return q * stepsize + baby_log_dict[now];
-            (now *= giant_a) %= M / g;
-        }
-        return -1;
+// Solve min_n f^n s = t (0 <= n <= max_search) f \in F, s \in S, t \in S
+// mapping: (F, S) -> S
+// composition: (F, F) -> F
+template <class S, class F, class Container>
+long long
+DiscreteLogarithm(const F &f, const S &s, const S &t, const std::function<S(F, S)> &mapping,
+                  const std::function<F(F, F)> &composition, long long max_search) {
+    const int giant_stride = ceil(sqrtl(max_search));
+    F giant = f, tmp = f;
+    for (int n = giant_stride - 1; n; n >>= 1) {
+        if (n & 1) giant = composition(giant, tmp);
+        tmp = composition(tmp, tmp);
     }
-};
+    return DiscreteLogarithm<S, F, Container>(f, giant, s, t, mapping, max_search, giant_stride);
+}
+
+// Solve min_n x^n = y (1 <= n <= max_search)
+template <class S, class Container>
+long long DiscreteLogarithmNonzero(const S &x, const S &y, const std::function<S(S, S)> &op,
+                                   long long max_search) {
+    long long res = DiscreteLogarithm<S, S, Container>(x, x, y, op, op, max_search);
+    if (res < 0 or res >= max_search) return -1;
+    return res + 1;
+}
+
+// Solve min_n x^n = y mod md (n >= 0) or return -1 if infeasible
+template <class Int> Int DiscreteLogarithmMod(Int x, Int y, Int md) {
+    x %= md, y %= md;
+    if (x < 0) x += md;
+    if (y < 0) y += md;
+    // You may change __int128 to long long, but be careful about overflow.
+    auto f = [&](Int a, Int b) -> Int { return __int128(a) * b % md; };
+    return DiscreteLogarithm<Int, Int, std::unordered_set<Int>>(x, Int{1} % md, y, f, f, md);
+}
+
+// Solve min_n x^n = y mod md (n >= 1) or return -1 if infeasible
+template <class Int> Int DiscreteLogarithmModNonzero(Int x, Int y, Int md) {
+    x %= md, y %= md;
+    if (x < 0) x += md;
+    if (y < 0) y += md;
+    // You may change __int128 to long long, but be careful about overflow.
+    auto f = [&](Int a, Int b) -> Int { return __int128(a) * b % md; };
+    return DiscreteLogarithmNonzero<Int, std::unordered_set<Int>>(x, y, f, md);
+}

number/discrete_logarithm.md

@@ -0,0 +1,96 @@
+---
+title: Discrete logarithm / Baby-step giant-step （離散対数問題）
+documentation_of: ./discrete_logarithm.hpp
+---
+
+以下， $S$ を集合， $F$ を $S$ から $S$ への写像の集合で，特に合成写像が効率的に計算できるものとする（ [AC Library の Lazy Segtree]([atcoder.github.io/ac-library/production/document_ja/lazysegtree.html](https://atcoder.github.io/ac-library/production/document_ja/lazysegtree.html)) とおおむね同様）．
+このコードは， $s, t \in S$ および $f \in F$ が与えられたとき， $f^n (s) = t, 0 \le n \le m$ を満たす最小の非負整数 $n$ を $O(\sqrt{m})$ で求める．
+
+ここで $S$ の位数が既知の場合， $m$ にはその値を指定すれば解が（存在すれば）必ず見つかるので， $S$ の位数があまり大きくない場合は高速に解が見つかることが期待できる．
+
+また， $S$ が特にモノイドのとき， $S$ の元 $x, y$ について， $x^n = y$ を満たす最小の正の $n$ を求めることもできる．
+
+## 実装されているアルゴリズムについて
+
+$F$ の元 $f$ を固定すると， $S$ の各元 $x$ から $f(x)$ に有向辺を張ってできる有向グラフは functional graph となる．
+
+$S$ の元 $s, t$ が与えられ， $f^n (s) = t$ を満たす $n$ を求めたい状況を考える．このコードは，適当な非負整数 $k$ をとって $t, f(t), f^2 (t), \ldots, f^{k - 1} (t)$ を予め計算しておき， $s$ に対して $f^k$ を繰り返し作用させてこれらのいずれかに一致するタイミングを検索する．
+
+なおこの検索は， 2 度マッチするまでループを回す必要がある点に注意（上述の functional graph 上で， $s$ を始点とするパスが $t$ を始点とするパスと $t$ から少しだけ進んだタイミングで合流し，その後ぐるっと回って $t$ に初めて到達するようなケースがあるので）．
+
+## 使用方法
+
+### 一般の問題（ $f^n (s) = t$ ）
+
+$f^n (s) = t$, $0 \le n \le m$ を満たす最小の整数 $n$ を以下のコードで求める．解がない場合は `-1` を返す． `DiscreteLogarithm()` の内部で`mapping()` は $O(\sqrt{m})$ 回， `composition()` は $O(\log m)$ 回呼ばれる．
+
+```cpp
+auto mapping = [&](F f, S x) -> S {
+    /* Implement */
+};
+auto composition = [&](F f, F g) -> F {
+    /* Implement */
+};
+
+F f;
+S s, t;
+long long m;
+
+long long ret = DiscreteLogarithm<S, F, std::unordered_set<S>>(f, s, t, mapping, composition, m);
+```
+
+なお，特に $S = F$ の場合には $x, y \in S$ に対して以下のコードで $x^n = y$ を満たす最小の **正の** 整数 $n$ を以下のコードで求められる．
+
+```cpp
+auto op = [&](S x, S y) -> S {};
+S x, y;
+long long m;
+long long ret = DiscreteLogarithmNonzero(x, y, op, m);
+```
+
+### 剰余類環の離散対数問題（ $x^n \equiv y \mod m$ ）
+
+与えられた非負整数 $x, y, m$ について， $x^n \equiv y \mod m$ を満たす最小の非負整数 $n$ を返すか、存在しない場合 $-1$ を返す．計算量は $O(\sqrt{m})$ である．
+なお内部で乗算オーバーフロー回避のため `__int128` を使用しているが， $m$ が $10^9$ 程度に収まるなら `long long` に書き換えると若干高速になる．
+
+```cpp
+int x, y, m;
+const int res = DiscreteLogarithmMod(x, y, m);
+```
+
+最小の **正の** 整数解を求めたい場合は代わりに `DiscreteLogarithmModNonzero` 関数を使うとよい．
+
+## 問題例
+
+`S`, `F` の型に応じて分類した．
+
+### 剰余環
+
+- [Library Checker: Discrete Logarithm](https://judge.yosupo.jp/problem/discrete_logarithm_mod)
+- [AtCoder Regular Contest 042 D - あまり](https://atcoder.jp/contests/arc042/tasks/arc042_d)
+- [No.551 夏休みの思い出（２） - yukicoder](https://yukicoder.me/problems/no/551)
+
+### $S, F$: 対称群
+
+- [No.101 ぐるぐる！あみだくじ！ - yukicoder](https://yukicoder.me/problems/no/101)
+- [No.261 ぐるぐるぐるぐる！あみだくじ！ - yukicoder](https://yukicoder.me/problems/no/261)
+
+### $S, F$: 有限体上の 2 次正方行列
+
+行列の det が初めて一致する時刻と，作用の行列の累乗が 1 になる周期をそれぞれ離散対数問題として求めて，その上で再度完全に行列が一致する時刻を求める必要がある．詳細は下記の問題に対する本リポジトリのテストコード参照．
+
+- [No.3170 [Cherry 7th Tune KY] Even if you could say "See you ..." - yukicoder](https://yukicoder.me/problems/no/3170)
+- [No.950 行列累乗 - yukicoder](https://yukicoder.me/problems/no/950)
+- [Fibonacci Mod Round #59 (Div. 2 only)](https://csacademy.com/contest/round-59/task/fibonacci-mod)
+- [L. FF?DLP - Monoxer Programming Contest for Engineers](https://mofecoder.com/contests/monoxercon_202508/tasks/monoxercon_202508_l)
+  - 複素数 $x + yi$ を行列 $\begin{pmatrix} x & -y \\ y & x\end{pmatrix}$ と置き換えて考えるとよい．
+
+### $S$: （位数があまり大きくない）環， $F$: $S$ 上の一次関数
+
+- [N - Paken Machine](https://atcoder.jp/contests/pakencamp-2022-day1/tasks/pakencamp_2022_day1_n)
+- [AtCoder Beginner Contest 222 G - 222](https://atcoder.jp/contests/abc222/tasks/abc222_g)
+- [AtCoder Beginner Contest 270 G - Sequence in mod P](https://atcoder.jp/contests/abc270/tasks/abc270_g)
+
+## 参考文献・リンク
+
+- [モノイド作用に関する離散対数問題 | maspyのHP](https://maspypy.com/%e3%83%a2%e3%83%8e%e3%82%a4%e3%83%89%e4%bd%9c%e7%94%a8%e3%81%ab%e9%96%a2%e3%81%99%e3%82%8b%e9%9b%a2%e6%95%a3%e5%af%be%e6%95%b0%e5%95%8f%e9%a1%8c)

number/test/discrete_logarithm_matrix.yuki3170.test.cpp

@@ -0,0 +1,32 @@
+#define PROBLEM "https://yukicoder.me/problems/no/3170"
+#include "number/discrete_logarithm.hpp"
+
+#include <array>
+#include <iostream>
+#include <set>
+using namespace std;
+
+int main() {
+    cin.tie(nullptr), ios::sync_with_stdio(false);
+    int T;
+    cin >> T;
+    while (T--) {
+        using S = array<int, 4>;
+
+        int P;
+        cin >> P;
+
+        auto op = [&](S l, S r) -> S {
+            return S{(int)(((long long)l[0b00] * r[0b00] + (long long)l[0b01] * r[0b10]) % P),
+                     (int)(((long long)l[0b00] * r[0b01] + (long long)l[0b01] * r[0b11]) % P),
+                     (int)(((long long)l[0b10] * r[0b00] + (long long)l[0b11] * r[0b10]) % P),
+                     (int)(((long long)l[0b10] * r[0b01] + (long long)l[0b11] * r[0b11]) % P)};
+        };
+        array<int, 4> A, B;
+        for (auto &x : A) cin >> x;
+        for (auto &x : B) cin >> x;
+
+        auto res = DiscreteLogarithm<S, S, set<S>>(A, S{1, 0, 0, 1}, B, op, op, (long long)P * P);
+        cout << res << '\n';
+    }
+}

number/test/discrete_logarithm_matrix.yuki950.test.cpp

@@ -0,0 +1,65 @@
+#define PROBLEM "https://yukicoder.me/problems/no/950"
+#include "number/discrete_logarithm.hpp"
+#include "utilities/pow_op.hpp"
+
+#include <array>
+#include <iostream>
+#include <set>
+using namespace std;
+
+using S = array<long long, 4>;
+
+long long P;
+S op(S l, S r) {
+    return S{(((long long)l[0b00] * r[0b00] % P + (long long)l[0b01] * r[0b10] % P) % P),
+             (((long long)l[0b00] * r[0b01] % P + (long long)l[0b01] * r[0b11] % P) % P),
+             (((long long)l[0b10] * r[0b00] % P + (long long)l[0b11] * r[0b10] % P) % P),
+             (((long long)l[0b10] * r[0b01] % P + (long long)l[0b11] * r[0b11] % P) % P)};
+};
+
+long long det(S x) { return ((long long)x[0] * x[3] % P - (long long)x[1] * x[2] % P + P) % P; };
+
+int main() {
+    cin.tie(nullptr), ios::sync_with_stdio(false);
+    cin >> P;
+
+    S A, B;
+    for (auto &x : A) cin >> x;
+    for (auto &x : B) cin >> x;
+
+    const long long det_A = det(A), det_B = det(B);
+
+    if (!det_A and det_B) {
+        puts("-1");
+        return 0;
+    }
+
+    const long long first_det_match = DiscreteLogarithmModNonzero<long long>(det_A, det_B, P);
+
+    if (first_det_match < 0) {
+        puts("-1");
+        return 0;
+    }
+
+    if (pow_op<S>(A, first_det_match, op) == B) {
+        cout << first_det_match << '\n';
+        return 0;
+    }
+
+    const long long det_period = det_A ? DiscreteLogarithmModNonzero<long long>(det_A, 1, P) : 1;
+    if (det_period < 0) {
+        puts("-1");
+        return 0;
+    }
+
+    const S init = pow_op<S>(A, first_det_match, op);
+    const S x = pow_op<S>(A, det_period, op);
+
+    const long long res = DiscreteLogarithm<S, S, set<S>>(x, init, B, op, op, P * 2);
+
+    if (res < 0) {
+        puts("-1");
+    } else {
+        cout << first_det_match + res * det_period << '\n';
+    }
+}

number/test/discrete_logarithm_mod.test.cpp

@@ -0,0 +1,19 @@
+#define PROBLEM "https://judge.yosupo.jp/problem/discrete_logarithm_mod"
+#include "number/discrete_logarithm.hpp"
+
+#include <iostream>
+using namespace std;
+
+int X, Y, M;
+
+int mp(int x, int y) { return (long long)x * y % M; }
+
+int main() {
+    cin.tie(nullptr), ios::sync_with_stdio(false);
+    int T;
+    cin >> T;
+    while (T--) {
+        cin >> X >> Y >> M;
+        cout << DiscreteLogarithmMod<long long>(X, Y, M) << '\n';
+    }
+}