123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225 |
- #!/usr/bin/ruby
- # Daniel "Trizen" Șuteu
- # Date: 26 October 2017
- # https://github.com/trizen
- # A recursive algorithm for finding all the non-negative integer solutions to the equation:
- # a^2 + b^2 = n
- # for any given positive integer `n` for which such a solution exists.
- # Example:
- # 99025 = 41^2 + 312^2 = 48^2 + 311^2 = 95^2 + 300^2 = 104^2 + 297^2 = 183^2 + 256^2 = 220^2 + 225^2
- # Blog post:
- # https://trizenx.blogspot.com/2017/10/representing-integers-as-sum-of-two.html
- # This algorithm is efficient when the factorization of `n` can be computed.
- # See also:
- # https://oeis.org/A001481
- func tonelli(n, p) {
- var q = p-1
- var s = valuation(q, 2)
- s == 1 && return powmod(n, (p + 1) >> 2, p)
- q >>= s
- var c = powmod(2 ..^ p -> first {|z| legendre(z, p) == -1 }, q, p)
- var r = powmod(n, (q + 1) >> 1, p)
- var t = powmod(n, q, p)
- while (!p.divides(t - 1)) {
- var b = 1
- var t2 = (t*t % p)
- for i in (1 ..^ s) {
- if (p.divides(t2 - 1)) {
- b = powmod(c, 1 << (s - i - 1), p)
- s = i
- break
- }
- t2 = (t2*t2 % p)
- }
- r = (r*b % p)
- c = (b*b % p)
- t = (t*c % p)
- }
- return r
- }
- func sqrt_mod_n(a, n) is cached {
- a = (a % n)
- kronecker(a, n) == 1 || return []
- if ((n & (n - 1)) == 0) { # n is a power of 2
- if (a % 8 == 1) {
- var k = n.valuation(2)
- k == 1 && return [1]
- k == 2 && return [1, 3]
- k == 3 && return [1, 3, 5, 7]
- if (a == 1) {
- return [1, (n>>1) - 1, (n>>1) + 1, n - 1]
- }
- return gather {
- for s in (sqrt_mod_n(a, n >> 1)) {
- var i = (((s*s - a) >> (k - 1)) % 2)
- var r = (s + (i << (k - 2)))
- take(r, n - r)
- }
- }.uniq
- }
- return []
- }
- if (n.is_prime) { # n is a prime
- return gather {
- var r = tonelli(a, n)
- take(r, n - r)
- }
- }
- var pe = n.factor_exp # factorize `n` into prime powers
- if (pe.len == 1) { # `n` is an odd prime power
- var p = pe[0][0]
- var k = pe[0][1]
- kronecker(a, p) == 1 || return []
- var (r1, r2) = with( tonelli(a, p) ) { |r|
- (r, n - r)
- }
- var pk = p
- var pi = p*p
- (k-1).times {
- var x = r1
- var y = (invmod(2, pk) * invmod(x, pk))
- r1 = ((pi + x - y*(x*x - a + pi)) % pi)
- r2 = (pi - r1)
- pk *= p
- pi *= p
- }
- return [r1, r2]
- }
- var solutions = []
- for p,e in (pe) {
- var m = p**e
- var r = sqrt_mod_n(a, m)
- solutions << r.map {|r0| [r0, m] }
- }
- gather {
- solutions.cartesian {|*a|
- take(Math.chinese(a...))
- }
- }.uniq
- }
- func sum_of_two_squares_solutions(n) is cached {
- n < 0 && return []
- n == 0 && return [[0, 0]]
- var prod1 = 1
- var prod2 = 1
- var prime_powers = []
- for p,e in (n.factor_exp) {
- if (p % 4 == 3) { # p = 3 (mod 4)
- e.is_even || return [] # power must be even
- prod2 *= p**(e >> 1)
- }
- elsif (p == 2) { # p = 2
- if (e.is_even) { # power is even
- prod2 *= p**(e >> 1)
- }
- else { # power is odd
- prod1 *= p
- prod2 *= p**((e - 1) >> 1)
- prime_powers << [p, 1]
- }
- }
- else { # p = 1 (mod 4)
- prod1 *= p**e
- prime_powers << [p, e]
- }
- }
- prod1 == 1 && return [[prod2, 0]]
- prod1 == 2 && return [[prod2, prod2]]
- var solutions = []
- for r in (sqrt_mod_n(-1, prod1)) {
- var s = r
- var q = prod1
- while (s*s > prod1) {
- (s, q) = (q % s, s)
- }
- solutions << [prod2 * s, prod2 * (q % s)]
- }
- for p,e in (prime_powers) {
- for (var i = e%2; i < e; i += 2) {
- var sq = p**((e - i) >> 1)
- var pp = p**(e - i)
- solutions += __FUNC__(prod1 / pp).map { |pair|
- pair.map {|r| sq * prod2 * r }
- }
- }
- }
- solutions.map {|pair| pair.sort } \
- .uniq_by {|pair| pair[0] } \
- .sort_by {|pair| pair[0] }
- }
- 50.times {
- var n = 1e10.irand
- var solutions = sum_of_two_squares_solutions(n) || next
- say %Q(#{n} = #{solutions.map {|a| "#{a[0]}^2 + #{a[1]}^2" }.join(' = ') })
- }
- assert_eq(sum_of_two_squares_solutions(2025), [[0, 45], [27, 36]])
- assert_eq(sum_of_two_squares_solutions(164025), [[0, 405], [243, 324]])
- assert_eq(sum_of_two_squares_solutions(99025), [[41, 312], [48, 311], [95, 300], [104, 297], [183, 256], [220, 225]])
- assert_eq(
- -10 .. 160 -> grep { sum_of_two_squares_solutions(_).len > 0 },
- %n[0, 1, 2, 4, 5, 8, 9, 10, 13, 16, 17, 18, 20, 25, 26, 29, 32, 34, 36, 37, 40, 41, 45, 49, 50, 52, 53, 58, 61, 64, 65, 68, 72, 73, 74, 80, 81, 82, 85, 89, 90, 97, 98, 100, 101, 104, 106, 109, 113, 116, 117, 121, 122, 125, 128, 130, 136, 137, 144, 145, 146, 148, 149, 153, 157, 160]
- )
- assert_eq(
- sum_of_two_squares_solutions(11392163240756069707031250),
- [[39309472125, 3374998963875], [216763660575, 3368260197225], [477329304375, 3341305130625], [729359177085, 3295481517405], [735019741071, 3294223614297], [907262616645, 3251005657515], [982736803125, 3228992353125], [1151205969375, 3172835964375], [1224793301193, 3145162095999], [1393801568775, 3074000720175], [1622919634875, 2959441687125], [1847545189875, 2824666354125], [1993551800625, 2723584854375], [2056446956025, 2676413487825], [2194367046795, 2564549961435], [2198769707673, 2560776252111], [2386646521875, 2386646521875]]
- )
|