feat: Add primitive root implementation and corresponding test cases

Author: weilycoder

test/primitive_root.test.cpp

@@ -0,0 +1,19 @@
+#define PROBLEM "https://judge.yosupo.jp/problem/primitive_root"
+
+#include "../weilycoder/number-theory/primitive_root.hpp"
+#include <iostream>
+using namespace std;
+using namespace weilycoder;
+
+int main() {
+  cin.tie(nullptr)->sync_with_stdio(false);
+  cin.exceptions(cin.failbit | cin.badbit);
+  size_t t;
+  cin >> t;
+  while (t--) {
+    uint64_t p;
+    cin >> p;
+    cout << prime_primitive_root<>(p) << '\n';
+  }
+  return 0;
+}

weilycoder/number-theory/primitive_root.hpp

@@ -0,0 +1,73 @@
+#ifndef WEILYCODER_NUMBER_THEORY_PRIMITIVE_ROOT_HPP
+#define WEILYCODER_NUMBER_THEORY_PRIMITIVE_ROOT_HPP
+
+#include "factorize.hpp"
+#include "mod_utility.hpp"
+#include <array>
+#include <cstdint>
+#include <vector>
+
+namespace weilycoder {
+/**
+ * @brief Check if g is a primitive root modulo p
+ * @tparam N The size of the factors array
+ * @tparam bit32 Whether to use 32-bit modular multiplication
+ * @param g The candidate primitive root
+ * @param p The prime modulus
+ * @param factors The prime factors of p - 1
+ * @return true if g is a primitive root modulo p, false otherwise
+ */
+template <size_t N = 64, bool bit32 = false>
+constexpr bool is_primitive_root(uint64_t g, uint64_t p,
+                                 const std::array<uint64_t, N> &factors) {
+  for (size_t i = 0; i < N; ++i) {
+    uint64_t q = factors[i];
+    if (q == 0)
+      break;
+    if (fast_power_64<bit32>(g, (p - 1) / q, p) == 1)
+      return false;
+  }
+  return true;
+}
+
+/**
+ * @brief Check if g is a primitive root modulo p
+ * @tparam bit32 Whether to use 32-bit modular multiplication
+ * @param g The candidate primitive root
+ * @param p The prime modulus
+ * @param factors The prime factors of p - 1
+ * @return true if g is a primitive root modulo p, false otherwise
+ */
+template <bool bit32 = false>
+bool is_primitive_root(uint64_t g, uint64_t p, const std::vector<uint64_t> &factors) {
+  const size_t N = factors.size();
+  for (size_t i = 0; i < N; ++i) {
+    uint64_t q = factors[i];
+    if (q == 0)
+      break;
+    if (fast_power_64<bit32>(g, (p - 1) / q, p) == 1)
+      return false;
+  }
+  return true;
+}
+
+/**
+ * @brief Find a primitive root modulo a prime p
+ * @tparam bit32 Whether to use 32-bit modular multiplication
+ * @param p The prime modulus
+ * @return A primitive root modulo p, or 0 if p is not prime
+ */
+template <bool bit32 = false> constexpr uint64_t prime_primitive_root(uint64_t p) {
+  if (p == 2)
+    return 1;
+  if (!is_prime(p))
+    return 0;
+  auto factors = factorize_fixed<64, bit32>(p - 1);
+  for (uint64_t g = 2; g < p; ++g)
+    if (is_primitive_root<64, bit32>(g, p, factors))
+      return g;
+  return 0;
+}
+} // namespace weilycoder
+
+#endif