aboutsummaryrefslogtreecommitdiff
path: root/zmodn.h
blob: a24f2936691e47a08de52d05cd3787185f511c47 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
#ifndef ZMODN_H
#define ZMODN_H

#include <cstdint>
#include <iostream>
#include <optional>
#include <tuple>
#include <type_traits>

template<typename T>
concept Integer = requires(T a, T b, int i, std::ostream& os) {
	{T(i)};

	{a + b} -> std::same_as<T>;
	{a - b} -> std::same_as<T>;
	{a * b} -> std::same_as<T>;
	{a / b} -> std::same_as<T>;
	{a % b} -> std::same_as<T>;

	{a == b} -> std::same_as<bool>;
	{a != b} -> std::same_as<bool>;

	{os << a} -> std::same_as<std::ostream&>;
};

template<Integer T>
std::tuple<T, T, T> extended_gcd(T a, T b) {
	if (b == 0) return {a, 1, 0};
	auto [g, x, y] = extended_gcd(b, a%b);
	return {g, y, x - y*(a/b)};
}

template<Integer auto N>
requires(N > 1)
class Zmod {
public:
	Zmod(decltype(N) z) : value{(z%N + N) % N} {}
	decltype(N) toint() const { return value; }

	Zmod operator+(const Zmod& z) const { return value + z.value; }
	Zmod operator-(const Zmod& z) const { return value - z.value; }
	Zmod operator*(const Zmod& z) const { return value * z.value; }

	Zmod operator+=(const Zmod& z) { return (*this) = value + z.value; }
	Zmod operator-=(const Zmod& z) { return (*this) = value - z.value; }
	Zmod operator*=(const Zmod& z) { return (*this) = value * z.value; }

	Zmod operator^(decltype(N) z) const {
		if (z == 0) return 1;
		if (z % 2 == 0) return (((*this) * (*this)) ^ (z/2));
		return (*this) * ((*this) ^ (z-1));
	}

	bool operator==(const Zmod& z) const { return value == z.value; }
	bool operator!=(const Zmod& z) const { return value != z.value; }

	std::optional<Zmod> inverse() const {
		auto [g, a, _] = extended_gcd(value, N);
		return g == 1 ? Zmod(a) : std::optional<Zmod>{};
	}

	std::optional<Zmod> operator/(const Zmod& d) const {
		auto i = d.inverse();
		return i ? (*this) * i.value() : i;
	}

	std::optional<Zmod> operator/=(const Zmod& d) {
		auto q = *this / d;
		return q ? (*this = q.value()) : q;
	}

	friend std::ostream& operator<<(std::ostream& os, const Zmod<N>& z) {
		return os << "(" << z.value << " mod " << N << ")";
	}
private:
	decltype(N) value;
};

#endif