diff --git a/cmd/quadraticResidues.go b/cmd/quadraticResidues.go index ad122e8..f36d260 100644 --- a/cmd/quadraticResidues.go +++ b/cmd/quadraticResidues.go @@ -26,12 +26,13 @@ import ( ) var quadraticResiduesN uint +var quadraticResiduesCoprime bool func quadraticResidues(cmd *cobra.Command, args []string) { bufStdout := bufio.NewWriter(os.Stdout) defer bufStdout.Flush() - ch := sieve.QuadraticResidues(quadraticResiduesN, 1000) + ch := sieve.QuadraticResidues(quadraticResiduesN, quadraticResiduesCoprime, 1000) for i := 0; ; i++ { v, ok := <-ch if !ok { @@ -69,4 +70,6 @@ func init() { quadraticResiduesCmd.Flags().UintVarP(&quadraticResiduesN, "limit", "n", 0, "upper limit") quadraticResiduesCmd.MarkFlagRequired("limit") + + quadraticResiduesCmd.Flags().BoolVarP(&quadraticResiduesCoprime, "coprime-only", "c", false, "only count residues coprime to the modulus") } diff --git a/internal/lib/sieve/qresidue.go b/internal/lib/sieve/qresidue.go index a38b156..b09840d 100644 --- a/internal/lib/sieve/qresidue.go +++ b/internal/lib/sieve/qresidue.go @@ -16,22 +16,61 @@ along with this program. If not, see . */ package sieve -func updatePowersOfTwo(sieve []uint, n uint) { +// NOTE: these formulas come from https://web.archive.org/web/20151224013638/http://www.maa.org/sites/default/files/Walter_D22068._Stangl.pdf + +func updatePowersOfTwoCoprime(sieve []uint, n uint) { for q := uint(8); 2*q < n; q *= 2 { + // q(2^n) = 2^(n-3) for n >= 3 sieve[2*q] = 2 * sieve[q] } } -func updatePowersOfOddPrimes(sieve []uint, p uint, n uint) { +func updatePowersOfOddPrimesCoprime(sieve []uint, p uint, n uint) { for q := p; p*q < n; q *= p { + // q(p^n) = (p^n - p^(n-1)) / 2 sieve[p*q] = (p*q - q) / 2 } } +func updatePowersOfTwo(sieve []uint, n uint) { + k := 1 + for q := uint(1); 2*q < n; q *= 2 { + if k%2 == 0 { + // s(2^n) = (2^(n-1) + 4) / 3 for even n + sieve[2*q] = (q + 4) / 3 + } else { + // s(2^n) = (2^(n-1) + 5) / 3 for odd n + sieve[2*q] = (q + 5) / 3 + } + + k += 1 + } +} + +func updatePowersOfOddPrimes(sieve []uint, p uint, n uint) { + k := 2 + for q := p; p*q < n; q *= p { + if p == q { + // s(p^2) = (p^2 - p + 2) / 2 + sieve[p*q] = (p*p - p + 2) / 2 + } else if k%2 == 0 { + // s(p^n) = (p^(n+1) + p + 2) / (2*(p+1)) for even n + sieve[p*q] = (p*p*q + p + 2) / (2 * (p + 1)) + } else { + // s(p^n) = (p^(n+1) + 2*p + 1) / (2*(p+1)) for odd n + sieve[p*q] = (p*p*q + 2*p + 1) / (2 * (p + 1)) + } + + k += 1 + } +} + /* QuadraticResidues computes the number of quadratic residues modulo k for k=1 to n. + +see https://oeis.org/A046073 */ -func QuadraticResidues(n uint, buflen uint) chan uint { +func QuadraticResidues(n uint, coprime bool, buflen uint) chan uint { sieve := make([]uint, n) for i := uint(0); i < n; i++ { sieve[i] = 1 @@ -40,23 +79,50 @@ func QuadraticResidues(n uint, buflen uint) chan uint { ch := make(chan uint, buflen) go func() { defer close(ch) - for i := uint(0); i < n; i++ { - if i == 0 || i == 1 || i == 4 || i == 6 || i == 8 || i == 12 || i == 24 || sieve[i] != 1 { - ch <- sieve[i] - continue - } - - if i == 2 { - updatePowersOfTwo(sieve, n) - } else { - sieve[i] = (i - 1) / 2 - updatePowersOfOddPrimes(sieve, i, n) - } - - updateMultiples(sieve, i, n, false) - ch <- sieve[i] + if coprime { + sieveQRCoprime(sieve, n, ch) + } else { + sieveQR(sieve, n, ch) } }() return ch } + +func sieveQRCoprime(sieve []uint, n uint, ch chan uint) { + for i := uint(0); i < n; i++ { + if i == 0 || i == 1 || i == 4 || i == 6 || i == 8 || i == 12 || i == 24 || sieve[i] != 1 { + ch <- sieve[i] + continue + } + + if i == 2 { + updatePowersOfTwoCoprime(sieve, n) + } else { + sieve[i] = (i - 1) / 2 + updatePowersOfOddPrimesCoprime(sieve, i, n) + } + + updateMultiples(sieve, i, n, false) + ch <- sieve[i] + } +} + +func sieveQR(sieve []uint, n uint, ch chan uint) { + for i := uint(0); i < n; i++ { + if i == 0 || i == 1 || sieve[i] != 1 { + ch <- sieve[i] + continue + } + + if i == 2 { + updatePowersOfTwo(sieve, n) + } else { + sieve[i] = (i + 1) / 2 + updatePowersOfOddPrimes(sieve, i, n) + } + + updateMultiples(sieve, i, n, false) + ch <- sieve[i] + } +}