diff --git a/cmd/divisors.go b/cmd/divisors.go index 5ca2994..4d57e05 100644 --- a/cmd/divisors.go +++ b/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") }