perl-scripts/Math/sum_of_two_squares_all_solutions.pl

#!/usr/bin/perl

# Daniel "Trizen" Șuteu
# Date: 25 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

# This algorithm is efficient when the factorization of `n` can be computed.

# Blog post:
#   https://trizenx.blogspot.com/2017/10/representing-integers-as-sum-of-two.html

# See also:
#   https://oeis.org/A001481

use 5.020;
use strict;
use warnings;

use experimental qw(signatures);

use Math::GMPz;
use ntheory qw(sqrtmod factor_exp chinese forsetproduct);

sub sum_of_two_squares_solutions ($n) {

    $n < 0  and return;
    $n == 0 and return [0, 0];

    my %sqrtmod_cache;

    my $find_solutions = sub ($factor_exp) {

        my $prod1 = 1;
        my $prod2 = 1;

        my @prod1_factor_exp;

        foreach my $f (@$factor_exp) {
            my ($p, $e) = @$f;
            if ($p % 4 == 3) {    # p = 3 (mod 4)
                $e % 2 == 0 or return;    # power must be even
                $prod2 *= Math::GMPz->new($p)**($e >> 1);
            }
            elsif ($p == 2) {             # p = 2
                if ($e % 2 == 0) {        # power is even
                    $prod2 *= Math::GMPz->new($p)**($e >> 1);
                }
                else {                    # power is odd
                    $prod1 *= $p;
                    $prod2 *= Math::GMPz->new($p)**(($e - 1) >> 1);
                    push @prod1_factor_exp, [$p, 1];
                }
            }
            else {                        # p = 1 (mod 4)
                $prod1 *= Math::GMPz->new($p)**$e;
                push @prod1_factor_exp, $f;
            }
        }

        $prod1 == 1 and return [$prod2, 0];
        $prod1 == 2 and return [$prod2, $prod2];

        my @congruences;

        foreach my $pe (@prod1_factor_exp) {
            my ($p, $e) = @$pe;
            my $pp  = Math::GMPz->new($p)**$e;
            my $key = Math::GMPz::Rmpz_get_str($pp, 10);
            my $r = (
                $sqrtmod_cache{$key} //= sqrtmod(-1, $pp) // do {
                    require Math::Sidef;
                    Math::Sidef::sqrtmod(-1, $pp);
                }
            );
            $r = Math::GMPz->new("$r") if ref($r);
            push @congruences, [[$r, $pp], [$pp - $r, $pp]];
        }

        my @square_roots;

        forsetproduct {
            push @square_roots, Math::GMPz->new(chinese(@_));
        } @congruences;

        my @solutions;

        foreach my $r (@square_roots) {

            my $s = $r;
            my $q = $prod1;

            while ($s * $s > $prod1) {
                ($s, $q) = ($q % $s, $s);
            }

            push @solutions, [$prod2 * $s, $prod2 * ($q % $s)];
        }

        foreach my $pe (@prod1_factor_exp) {
            my ($p, $e) = @$pe;

            for (my $i = $e % 2 ; $i < $e ; $i += 2) {

                my @factor_exp;
                foreach my $pp (@prod1_factor_exp) {
                    if ($pp->[0] == $p) {
                        push(@factor_exp, [$p, $i]) if ($i > 0);
                    }
                    else {
                        push @factor_exp, $pp;
                    }
                }

                my $sq = $prod2 * Math::GMPz->new($p)**(($e - $i) >> 1);

                push @solutions, map {
                    [map { $_ * $sq } @$_]
                } __SUB__->(\@factor_exp);
            }
        }

        return @solutions;
    };

    my @factor_exp = factor_exp($n);
    my @solutions  = $find_solutions->(\@factor_exp);

    @solutions = sort { $a->[0] <=> $b->[0] } do {
        my %seen;
        grep { !$seen{$_->[0]}++ } map {
            [sort { $a <=> $b } @$_]
        } @solutions;
    };

    return @solutions;
}

# Run some tests

use Test::More tests => 6;

is_deeply([sum_of_two_squares_solutions(2025)],   [[0, 45],  [27,  36]],);
is_deeply([sum_of_two_squares_solutions(164025)], [[0, 405], [243, 324]]);
is_deeply([sum_of_two_squares_solutions(99025)],  [[41, 312], [48, 311], [95, 300], [104, 297], [183, 256], [220, 225]]);

is_deeply(
          [grep { my @arr = sum_of_two_squares_solutions($_); @arr > 0 } -10 .. 160],
          [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
          ]
         );

do {
    use bigint try => 'GMP';
    is_deeply(
              [sum_of_two_squares_solutions(Math::GMPz->new("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]
              ]
             );

    is_deeply([sum_of_two_squares_solutions(2**128 + 1)],
              [[1, 18446744073709551616], [8479443857936402504, 16382350221535464479]]);

    is_deeply(
              [sum_of_two_squares_solutions(13**18 * 5**7)],
              [[75291211970,   2963091274585],
               [100083884615,  2962357487570],
               [124869548830,  2961416259815],
               [149646468985,  2960267657230],
               [154416779750,  2960022656375],
               [179181003625,  2958626849750],
               [203932680250,  2957023863625],
               [228670076375,  2955213810250],
               [253391459750,  2953196816375],
               [258150241063,  2952784638466],
               [282850264814,  2950521038023],
               [307530481817,  2948050825694],
               [332189163826,  2945374174457],
               [356824584103,  2942491271746],
               [481345955350,  2924702504425],
               [505803171575,  2920572173350],
               [530224968650,  2916237327575],
               [554609636425,  2911698270650],
               [578955467350,  2906955320425],
               [583639307225,  2906018552450],
               [607936593550,  2901032879225],
               [632191308775,  2895844059550],
               [656401754450,  2890452456775],
               [680566235225,  2884858448450],
               [802350873038,  2853386013959],
               [826200069721,  2846571993278],
               [849991411282,  2839558639801],
               [873723231719,  2832346444642],
               [897393869198,  2824935912839],
               [901945120375,  2823486084250],
               [925540625750,  2815839700375],
               [949071319625,  2807996135750],
               [972535554250,  2799955939625],
               [977046452345,  2798385051790],
               [1000429281410, 2790111094745],
               [1023742054855, 2781641758610],
               [1046983140190, 2772977636455],
               [1070150909945, 2764119334990],
               [1186462080890, 2716226499895],
               [1209150070505, 2706203018090],
               [1231753388710, 2695990032905],
               [1254270452695, 2685588259510],
               [1276699685690, 2674998426295],
               [1281008818375, 2672937536750],
               [1303331253250, 2662124398375],
               [1325562421625, 2651124843250],
               [1347700766750, 2639939641625],
               [1369744738375, 2628569576750],
               [1373978929622, 2626358804329],
               [1395908335991, 2614769317862],
               [1417739993098, 2602996730711],
               [1439472372169, 2591041867258],
               [1461103951382, 2578905564649],
               [1569204922025, 2514592328950],
               [1590192225050, 2501373094025],
               [1611068173975, 2487978699050],
               [1631831306950, 2474410081975],
               [1652480170025, 2460668192950],
               [1656443419150, 2458002007175],
               [1676954116825, 2444054737150],
               [1697347384850, 2429936320825],
               [1717621795175, 2415647746850],
               [1838087734327, 2325298292486],
               [1857481600234, 2309835659287],
               [1876745394953, 2294211278554],
               [1895877769526, 2278426244393],
               [1899547017625, 2275368057250],
               [1918520912750, 2259392877625],
               [1937360462375, 2243259482750],
               [1956064347250, 2226969002375],
               [1974631257625, 2210522577250],
               [1978190975930, 2207337566135],
               [1996592834665, 2190706671530],
               [2014854880870, 2173922371465],
               [2032975835735, 2156985841270],
               [2050954430330, 2139898266935]
              ]
             )
      if 0;
};

my @nums = (@ARGV ? (map { Math::GMPz->new($_) } @ARGV) : (map { int rand(~0) } 1 .. 20));

foreach my $n (@nums) {
    (my @solutions = sum_of_two_squares_solutions($n)) || next;

    say "$n = " . join(' = ', map { "$_->[0]^2 + $_->[1]^2" } @solutions);

    # Verify solutions
    foreach my $solution (@solutions) {
        if ($n != $solution->[0]**2 + $solution->[1]**2) {
            die "error for $n: (@$solution)\n";
        }
    }
}

__END__
999826 = 99^2 + 995^2 = 315^2 + 949^2 = 525^2 + 851^2 = 699^2 + 715^2
999828 = 318^2 + 948^2
999844 = 312^2 + 950^2 = 410^2 + 912^2
999848 = 62^2 + 998^2
999850 = 43^2 + 999^2 = 321^2 + 947^2 = 565^2 + 825^2
999853 = 387^2 + 922^2
999857 = 401^2 + 916^2 = 544^2 + 839^2
999860 = 154^2 + 988^2 = 698^2 + 716^2
999869 = 262^2 + 965^2 = 613^2 + 790^2
999881 = 341^2 + 940^2 = 484^2 + 875^2
999882 = 309^2 + 951^2 = 651^2 + 759^2
999890 = 421^2 + 907^2 = 473^2 + 881^2
999892 = 324^2 + 946^2
999898 = 213^2 + 977^2 = 697^2 + 717^2
999909 = 222^2 + 975^2 = 678^2 + 735^2
999914 = 667^2 + 745^2
999917 = 109^2 + 994^2
999937 = 44^2 + 999^2 = 89^2 + 996^2
999938 = 77^2 + 997^2
999940 = 126^2 + 992^2 = 178^2 + 984^2 = 306^2 + 952^2 = 448^2 + 894^2 = 578^2 + 816^2 = 696^2 + 718^2
999941 = 370^2 + 929^2 = 446^2 + 895^2
999944 = 638^2 + 770^2
999946 = 585^2 + 811^2
999949 = 243^2 + 970^2 = 290^2 + 957^2 = 450^2 + 893^2 = 493^2 + 870^2
999952 = 444^2 + 896^2
999953 = 568^2 + 823^2
999954 = 327^2 + 945^2 = 375^2 + 927^2
999956 = 500^2 + 866^2
999961 = 644^2 + 765^2
999962 = 239^2 + 971^2 = 541^2 + 841^2
999968 = 452^2 + 892^2
999970 = 247^2 + 969^2 = 627^2 + 779^2
999973 = 63^2 + 998^2 = 118^2 + 993^2 = 273^2 + 962^2 = 442^2 + 897^2 = 622^2 + 783^2 = 658^2 + 753^2
999981 = 141^2 + 990^2
999986 = 365^2 + 931^2 = 695^2 + 719^2
999997 = 194^2 + 981^2 = 454^2 + 891^2
1000000 = 0^2 + 1000^2 = 280^2 + 960^2 = 352^2 + 936^2 = 600^2 + 800^2