123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115 |
- #!/usr/bin/ruby
- # Find all solutions to the quadratic congruence:
- # x^2 = a (mod n)
- # Based on algorithm by Hugo van der Sanden:
- # https://github.com/danaj/Math-Prime-Util/pull/55
- func sqrtmod_all(a, n) {
- n = -n if (n < 0)
- n == 0 && return []
- n == 1 && return [0]
- a = (a % n)
- func sqrtmod_pk(a, p, k) {
- var pk = p**k
- if (p `divides` a) {
- if (pk `divides` a) {
- var low = p**(k >> 1)
- var high = (k.is_odd ? (low*p) : low)
- return (^low -> map {|i| high * i })
- }
- var a2 = a/p
- return [] if !(p `divides` a2)
- var pj = pk/p
- return __FUNC__(a2/p, p, k-2).map {|q|
- ^p -> map {|i| q*p + i*pj }...
- }
- }
- var q = sqrtmod(a, pk)
- q.is_nan && return []
- return [q, pk-q] if (p != 2)
- return [q] if (k == 1)
- return [q, pk-q] if (k == 2)
- var pj = p**(k-1)
- var q2 = (q * (pj-1))%pk
- return [q, pk-q, q2, pk-q2]
- }
- var roots = []
- n.factor_map {|p,k|
- sqrtmod_pk(a, p, k).map {|r| [r, p**k] }
- }.cartesian {|*a|
- roots << Math.chinese(a...)
- }
- return roots.sort
- }
- say sqrtmod_all(1, 8) #=> [1, 3, 5, 7]
- say sqrtmod_all(120, 5045) #=> [1165, 3880]
- say sqrtmod_all(4095, 8469) #=> [1110, 1713, 3933, 4536, 6756, 7359]
- # Run some tests
- assert_eq(
- sqrtmod_all(-1, 13**18 * 5**7)
- %n(633398078861605286438568 2308322911594648160422943 6477255756527023177780182 8152180589260066051764557)
- )
- assert_eq(
- sqrtmod_all(2466, 5967),
- [120 237 426 543 1446 1563 1752 1869 2109 2226 2415 2532 3435 3552 3741 3858 4098 4215 4404 4521 5424 5541 5730 5847]
- )
- assert_eq(
- sqrtmod_all(7281, 9954),
- %n(1233 1611 1707 2085 4551 4929 5025 5403 7869 8247 8343 8721)
- )
- assert_eq(
- sqrtmod_all(1701, 6300),
- %n[399, 651, 1449, 1701, 2499, 2751, 3549, 3801, 4599, 4851, 5649, 5901]
- )
- assert_eq(
- sqrtmod_all(306, 810),
- %n[66, 96, 174, 204, 336, 366, 444, 474, 606, 636, 714, 744]
- )
- assert_eq(
- sqrtmod_all(2754, 6561),
- %n[126, 603, 855, 1332, 1584, 2061, 2313, 2790, 3042, 3519, 3771, 4248, 4500, 4977, 5229, 5706, 5958, 6435]
- )
- assert_eq(
- sqrtmod_all(17640, 48465),
- %n[2865, 7905, 8250, 13290, 19020, 24060, 24405, 29445, 35175, 40215, 40560, 45600]
- )
- for n in (0..20), a in (^n) {
- var roots = sqrtmod_all(a, n)
- assert(roots.all {|r|
- r**2 % n == a
- }, "sqrtmod(#{a}, #{n}) = #{roots}")
- assert_eq(
- ^n -> grep {|k| mulmod(k, k, n) == a },
- roots
- )
- }
|