add -e flag to set exponent

cmd/divisors.go

@@ -23,10 +23,26 @@ import (
 )
 
 var divisorsN uint
+var divisorsE uint
 
-func updateMultiples(sieve []uint, p uint, n uint) {
+func pow(base uint, exp uint) uint {
+	result := uint(1)
+	for exp > 0 {
+		if exp%2 == 1 {
+			result *= base
+		}
+
+		exp >>= 1
+		base *= base
+	}
+
+	return result
+}
+
+func updateMultiples(sieve []uint, x uint, p uint, n uint) {
 	q := p
 	for {
+		// sigma_x(a*b) = sigma_x(a) * sigma_x(b) if gcd(a,b) = 1
 		for i := 2 * q; i < n; i += q {
 			if i%(p*q) != 0 {
 				sieve[i] *= sieve[q]
@@ -37,12 +53,13 @@ func updateMultiples(sieve []uint, p uint, n uint) {
 			break
 		}
 
-		sieve[p*q] = p*q + sieve[q]
+		// sigma_x(p^k) = p^(kx) + sigma_x(p^(k-1))
+		sieve[p*q] = pow(p*q, x) + sieve[q]
 		q *= p
 	}
 }
 
-func divisorsSieve(n uint) chan uint {
+func divisorsSieve(n uint, x uint) chan uint {
 	sieve := make([]uint, n)
 	sieve[0] = 0
 	for i := uint(1); i < n; i++ {
@@ -57,8 +74,8 @@ func divisorsSieve(n uint) chan uint {
 				continue
 			}
 
-			sieve[i] = i + 1
-			updateMultiples(sieve, i, n)
+			sieve[i] = pow(i, x) + 1
+			updateMultiples(sieve, x, i, n)
 			ch <- sieve[i]
 		}
 
@@ -69,7 +86,7 @@ func divisorsSieve(n uint) chan uint {
 }
 
 func divisors(cmd *cobra.Command, args []string) {
-	ch := divisorsSieve(divisorsN)
+	ch := divisorsSieve(divisorsN, divisorsE)
 	for i := 0; ; i++ {
 		v, ok := <-ch
 		if !ok {
@@ -107,4 +124,6 @@ func init() {
 
 	divisorsCmd.Flags().UintVarP(&divisorsN, "limit", "n", 0, "upper limit")
 	divisorsCmd.MarkFlagRequired("limit")
+
+	divisorsCmd.Flags().UintVarP(&divisorsE, "exponent", "e", 0, "exponent")
 }