Skip to content

Instantly share code, notes, and snippets.

@phipsgabler
Last active February 28, 2018 10:36
Show Gist options
  • Save phipsgabler/7b245cdf840fd026f3dc24d93793a3dd to your computer and use it in GitHub Desktop.
Save phipsgabler/7b245cdf840fd026f3dc24d93793a3dd to your computer and use it in GitHub Desktop.
Optimizing implementation of sigma function
using Primes
using BenchmarkTools
function σ{T}(k, n, ::Type{T} = Int)
result = ones(T, n - 1)
for p in primes(n)
bound = round(Int, log(n) / log(p), RoundDown)
pᵏ = convert(T, p)^k
pᵏⁱ = one(T)
si = pᵏⁱ
pi = 1
for i in 0:bound
for pj in eachindex(result)
if pi * pj < n
@inbounds result[pi*pj] = si * result[pj]
else
break
end
end
pi *= p
pᵏⁱ *= pᵏ
si += pᵏⁱ
end
end
return result
end
# Original code from https://stackoverflow.com/q/49021937/1346276
thefunc(i) = (a::Int, b::Int)-> sum(BigInt(a)^(k*i) for k in 0:b)
function powsfunc(x, n, func)
bound = convert(Int, floor(log(n)/log(x)))
return Dict{Int, Number}(x^i => func(x, i) for i in 0:bound)
end
dictprod2(d1, d2, bd)=Dict(i*j=> d1[i]*d2[j] for i in keys(d1), j in keys(d2) if i*j<bd)
function dosigma(k, n)
primefunc = thefunc(k)
combfunc = (a, b) -> dictprod2(a, b, n)
allprimes = [powsfunc(p, n, primefunc) for p in primes(n)]
trivdict = Dict{Int, Number}(1=>1)
theresult = reduce(combfunc, trivdict, allprimes)
return theresult
end
function test(;k = 0, n = 100, T = Int)
a = @btime σ($k, $n, $T)
b = dosigma(k, n)
@assert all(a[i] == b[i] for i = 1:(n-1))
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment