add flag for whether to only count coprime residues

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")
 }

internal/lib/sieve/qresidue.go

@@ -16,22 +16,61 @@ along with this program. If not, see <http://www.gnu.org/licenses/>.
 */
 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]
+	}
+}