Skip to content

Commit 8384006

Browse files
Ryoga-exeRyoga-exe
authored
Merge pull request #13 from Ryoga-exe/feature/modint
Implement modint
2 parents 92d1b4f + aefe8f7 commit 8384006

File tree

3 files changed

+650
-0
lines changed

3 files changed

+650
-0
lines changed

‎src/internal_math.zig‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,174 @@ test Barrett {
5757
try std.testing.expectEqual(@as(u32, 2147483646), b3.mul(1073741824, 2147483645));
5858
}
5959

60+
// (x ^ n) % m
61+
pub fn comptimePowMod(x: comptime_int, n: comptime_int, m: comptime_int) comptime_int {
62+
if (!(0 <= n and 1 <= m)) {
63+
@compileError("Arguments do not satisfy 0 <= n and 1 <= m");
64+
}
65+
if (m == 1) {
66+
return 0;
67+
}
68+
var _n = n;
69+
var r = 1;
70+
var y = @mod(x, m);
71+
while (_n > 0) : (_n >>= 1) {
72+
if (_n & 1 == 1) {
73+
r = (r * y) % m;
74+
}
75+
y = (y * y) % m;
76+
}
77+
return r;
78+
}
79+
80+
test comptimePowMod {
81+
const i32max = std.math.maxInt(i32);
82+
const i64max = std.math.maxInt(i64);
83+
const tests = comptime &[_]struct { x: comptime_int, n: comptime_int, m: comptime_int, expects: comptime_int }{
84+
.{ .x = 0, .n = 0, .m = 1, .expects = 0 },
85+
.{ .x = 0, .n = 0, .m = 3, .expects = 1 },
86+
.{ .x = 0, .n = 0, .m = 723, .expects = 1 },
87+
.{ .x = 0, .n = 0, .m = 998244353, .expects = 1 },
88+
.{ .x = 0, .n = 0, .m = i32max, .expects = 1 },
89+
90+
.{ .x = 0, .n = 1, .m = 1, .expects = 0 },
91+
.{ .x = 0, .n = 1, .m = 3, .expects = 0 },
92+
.{ .x = 0, .n = 1, .m = 723, .expects = 0 },
93+
.{ .x = 0, .n = 1, .m = 998244353, .expects = 0 },
94+
.{ .x = 0, .n = 1, .m = i32max, .expects = 0 },
95+
96+
.{ .x = 0, .n = i64max, .m = 1, .expects = 0 },
97+
.{ .x = 0, .n = i64max, .m = 3, .expects = 0 },
98+
.{ .x = 0, .n = i64max, .m = 723, .expects = 0 },
99+
.{ .x = 0, .n = i64max, .m = 998244353, .expects = 0 },
100+
.{ .x = 0, .n = i64max, .m = i32max, .expects = 0 },
101+
102+
.{ .x = 1, .n = 0, .m = 1, .expects = 0 },
103+
.{ .x = 1, .n = 0, .m = 3, .expects = 1 },
104+
.{ .x = 1, .n = 0, .m = 723, .expects = 1 },
105+
.{ .x = 1, .n = 0, .m = 998244353, .expects = 1 },
106+
.{ .x = 1, .n = 0, .m = i32max, .expects = 1 },
107+
108+
.{ .x = 1, .n = 1, .m = 1, .expects = 0 },
109+
.{ .x = 1, .n = 1, .m = 3, .expects = 1 },
110+
.{ .x = 1, .n = 1, .m = 723, .expects = 1 },
111+
.{ .x = 1, .n = 1, .m = 998244353, .expects = 1 },
112+
.{ .x = 1, .n = 1, .m = i32max, .expects = 1 },
113+
114+
.{ .x = 1, .n = i64max, .m = 1, .expects = 0 },
115+
.{ .x = 1, .n = i64max, .m = 3, .expects = 1 },
116+
.{ .x = 1, .n = i64max, .m = 723, .expects = 1 },
117+
.{ .x = 1, .n = i64max, .m = 998244353, .expects = 1 },
118+
.{ .x = 1, .n = i64max, .m = i32max, .expects = 1 },
119+
120+
.{ .x = i64max, .n = 0, .m = 1, .expects = 0 },
121+
.{ .x = i64max, .n = 0, .m = 3, .expects = 1 },
122+
.{ .x = i64max, .n = 0, .m = 723, .expects = 1 },
123+
.{ .x = i64max, .n = 0, .m = 998244353, .expects = 1 },
124+
.{ .x = i64max, .n = 0, .m = i32max, .expects = 1 },
125+
126+
.{ .x = i64max, .n = i64max, .m = 1, .expects = 0 },
127+
.{ .x = i64max, .n = i64max, .m = 3, .expects = 1 },
128+
.{ .x = i64max, .n = i64max, .m = 723, .expects = 640 },
129+
.{ .x = i64max, .n = i64max, .m = 998244353, .expects = 683296792 },
130+
.{ .x = i64max, .n = i64max, .m = i32max, .expects = 1 },
131+
132+
.{ .x = 2, .n = 3, .m = 1_000_000_007, .expects = 8 },
133+
.{ .x = 5, .n = 7, .m = 1_000_000_007, .expects = 78125 },
134+
.{ .x = 123, .n = 456, .m = 1_000_000_007, .expects = 565291922 },
135+
};
136+
inline for (tests) |t| {
137+
try std.testing.expectEqual(t.expects, comptimePowMod(t.x, t.n, t.m));
138+
}
139+
}
140+
141+
// Reference:
142+
// M. Forisek and J. Jancina,
143+
// Fast Primality Testing for Integers That Fit into a Machine Word
144+
pub fn comptimeIsPrime(n: comptime_int) bool {
145+
if (n == 2) {
146+
return true;
147+
} else if (n <= 1 or n % 2 == 0) {
148+
return false;
149+
}
150+
const base = std.math.log2(n);
151+
const upper = (1 << base) - 1;
152+
const magnitude_bits = if (upper >= n) base else base + 1;
153+
154+
// Reference:
155+
// https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test
156+
// https://miller-rabin.appspot.com/
157+
// https://oeis.org/A014233
158+
const seq = switch (magnitude_bits) {
159+
0...32 => [_]comptime_int{ 2, 7, 61 },
160+
33...64 => [_]comptime_int{ 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37 },
161+
else => @compileError("n is too large"),
162+
};
163+
164+
inline for (seq) |v| {
165+
if (n == v) {
166+
return true;
167+
}
168+
}
169+
170+
const d = comptime blk: {
171+
var d = n - 1;
172+
while (d % 2 == 0) {
173+
d /= 2;
174+
}
175+
break :blk d;
176+
};
177+
inline for (seq) |a| {
178+
const t, const y = comptime blk: {
179+
var t = d;
180+
var y = comptimePowMod(a, t, n);
181+
while (t != n - 1 and y != 1 and y != n - 1) : (t <<= 1) {
182+
y = y * y % n;
183+
}
184+
break :blk .{ t, y };
185+
};
186+
if (y != n - 1 and t % 2 == 0) {
187+
return false;
188+
}
189+
}
190+
return true;
191+
}
192+
193+
test comptimeIsPrime {
194+
const tests = comptime &[_]struct { x: comptime_int, expects: bool }{
195+
.{ .x = 0, .expects = false },
196+
.{ .x = 1, .expects = false },
197+
.{ .x = 2, .expects = true },
198+
.{ .x = 3, .expects = true },
199+
.{ .x = 4, .expects = false },
200+
.{ .x = 5, .expects = true },
201+
.{ .x = 6, .expects = false },
202+
.{ .x = 7, .expects = true },
203+
.{ .x = 8, .expects = false },
204+
.{ .x = 9, .expects = false },
205+
206+
// .{ .x = 57, .expect = true },
207+
.{ .x = 57, .expects = false },
208+
.{ .x = 58, .expects = false },
209+
.{ .x = 59, .expects = true },
210+
.{ .x = 60, .expects = false },
211+
.{ .x = 61, .expects = true },
212+
.{ .x = 62, .expects = false },
213+
214+
.{ .x = 701928443, .expects = false },
215+
.{ .x = 998244353, .expects = true },
216+
.{ .x = 1_000_000_000, .expects = false },
217+
.{ .x = 1_000_000_007, .expects = true },
218+
219+
.{ .x = std.math.maxInt(i32), .expects = true },
220+
.{ .x = std.math.maxInt(i62), .expects = true },
221+
.{ .x = std.math.maxInt(i64), .expects = false },
222+
};
223+
inline for (tests) |t| {
224+
try std.testing.expectEqual(t.expects, comptimeIsPrime(t.x));
225+
}
226+
}
227+
60228
pub fn invGcd(a: i64, b: i64) struct { i64, i64 } {
61229
const c = @mod(a, b);
62230
if (c == 0) {

‎src/lib.zig‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,11 @@ pub const LazySegtreeNS = @import("lazysegtree.zig").LazySegtreeFromNS;
1111
pub const MfGraph = @import("maxflow.zig").MfGraph;
1212
pub const McfGraph = @import("mincostflow.zig").McfGraph;
1313
pub const TwoSat = @import("twosat.zig");
14+
pub const StaticModint = @import("modint.zig").StaticModint;
15+
pub const Modint1000000007 = @import("modint.zig").Modint1000000007;
16+
pub const Modint998244353 = @import("modint.zig").Modint998244353;
17+
pub const DynamicModint = @import("modint.zig").DynamicModint;
18+
pub const Modint = @import("modint.zig").Modint;
1419

1520
pub usingnamespace @import("math.zig");
1621
pub usingnamespace @import("string.zig");

0 commit comments

Comments
 (0)