Library/Math/BabyStepGiantStep.hpp

View this file on GitHub
Last update: 2025-01-09 01:07:21+09:00
Include: #include "Library/Math/BabyStepGiantStep.hpp"

Verified with

Test/Math/BabyStepGiantStep.test.cpp

Code

#pragma once

#include <cmath>
#include <optional>
#include <ranges>
#include <unordered_set>

namespace mtd {
  /*
   * x^i * s = t (0 <= i < n) を満たす最小のi
   * O(sqrt(n))
   */
  template <class Hash, class Monoid, class T, class Product, class Op>
  std::optional<long long> baby_step_giant_step(const Monoid& x, const T& s,
                                                const T& t, long long n,
                                                const Product prod,
                                                const Op& op) {
    if (n == 0) { return (s == t) ? std::make_optional(0) : std::nullopt; }
    if (n == 1) {
      return (s == t)          ? std::make_optional(0)
             : (op(x, s) == t) ? std::make_optional(1)
                               : std::nullopt;
    }

    auto m = static_cast<long long>(std::sqrt(n));
    std::unordered_set<T, Hash> mp;
    auto xm = x;
    for (auto i : std::views::iota(0, m)) {
      mp.emplace(op(xm, t));
      if (i < m - 1) { xm = prod(xm, x); }
    }

    bool find = false;
    auto xks = s;
    for (auto k : std::views::iota(1, n / m + 5)) {
      auto next = op(xm, xks);
      if (mp.contains(next)) {
        auto xi = xks;
        for (auto i : std::views::iota(m * (k - 1), m * k + 1)) {
          if (xi == t) { return i; }
          xi = op(x, xi);
        }
        if (find) { return std::nullopt; }
        find = true;
      }
      std::swap(xks, next);
    }
    return std::nullopt;
  }
};  // namespace mtd

#line 2 "Library/Math/BabyStepGiantStep.hpp"

#include <cmath>
#include <optional>
#include <ranges>
#include <unordered_set>

namespace mtd {
  /*
   * x^i * s = t (0 <= i < n) を満たす最小のi
   * O(sqrt(n))
   */
  template <class Hash, class Monoid, class T, class Product, class Op>
  std::optional<long long> baby_step_giant_step(const Monoid& x, const T& s,
                                                const T& t, long long n,
                                                const Product prod,
                                                const Op& op) {
    if (n == 0) { return (s == t) ? std::make_optional(0) : std::nullopt; }
    if (n == 1) {
      return (s == t)          ? std::make_optional(0)
             : (op(x, s) == t) ? std::make_optional(1)
                               : std::nullopt;
    }

    auto m = static_cast<long long>(std::sqrt(n));
    std::unordered_set<T, Hash> mp;
    auto xm = x;
    for (auto i : std::views::iota(0, m)) {
      mp.emplace(op(xm, t));
      if (i < m - 1) { xm = prod(xm, x); }
    }

    bool find = false;
    auto xks = s;
    for (auto k : std::views::iota(1, n / m + 5)) {
      auto next = op(xm, xks);
      if (mp.contains(next)) {
        auto xi = xks;
        for (auto i : std::views::iota(m * (k - 1), m * k + 1)) {
          if (xi == t) { return i; }
          xi = op(x, xi);
        }
        if (find) { return std::nullopt; }
        find = true;
      }
      std::swap(xks, next);
    }
    return std::nullopt;
  }
};  // namespace mtd