Mercurial > lbo > hg > juliaplay
changeset 28:743af1962848
arrays: correct nsum
author | Lewin Bormann <lbo@spheniscida.de> |
---|---|
date | Sat, 04 Mar 2023 16:39:37 +0100 |
parents | 668adaf0313f |
children | af82fbc3d6c0 |
files | julia/arrays.jl |
diffstat | 1 files changed, 31 insertions(+), 73 deletions(-) [+] |
line wrap: on
line diff
--- a/julia/arrays.jl Sat Mar 04 12:08:16 2023 +0100 +++ b/julia/arrays.jl Sat Mar 04 16:39:37 2023 +0100 @@ -47,6 +47,19 @@ return (0,0) end +function twosum_linear(a::AbstractVector{T}, s)::Vector{Pair} where {T<:Integer} + seen = Set{T}() + p = Pair[] + for e = a + f = s - e + if f in seen + push!(p, e => f) + end + push!(seen, e) + end + p +end + function twosum2(a::Vector{Int}, n::Int)::Tuple{Int,Int} # More clever? # @@ -76,83 +89,28 @@ nothing end -# About 11x faster than without memoization for input set of n = 11 -# MUCH MUCH faster for larger input sets -function threesum_memoize(a::AbstractVector{<:Integer}, s, i=1, running::Tuple=(), triplets=Set{Tuple}(), badstarts=Set{Tuple}()) - markbad(t) = if length(t) < 3 push!(badstarts, t) end - isbad(t) = t in badstarts - - if i > length(a) && length(running) < 3 - return 0 +function threesum(a::AbstractVector{T}, s) where {T<:Integer} + triples = Tuple[] + for e = a + f = s - e + pairs = twosum_linear(a, f) + append!(triples, ((e, p...) for p = pairs)) end - if length(running) == 3 - if sum(running) == s - push!(triplets, running) - return 1 - end - return 0 - elseif length(running) == 2 - current = a[i] - # incl/excl/new - possibilities = [running, (running..., current), (running[1], current), (running[2], current), ()] - for p = possibilities - if !isbad(p) && 0 == threesum_memoize(a, s, i+1, p, triplets, badstarts) - markbad(p) - end - end - elseif length(running) == 1 - current = a[i] - possibilities = [(running[1], current), running, ()] - # incl/excl current element - for p = possibilities - if !isbad(p) && 0 == threesum_memoize(a, s, i+1, p, triplets, badstarts) - markbad(p) - end - end - elseif length(running) == 0 - current = a[i] - if !isbad((current,)) && 0 == threesum_memoize(a, s, i+1, (current,), triplets, badstarts) - markbad((current,)) - end - if !isbad(running) && 0 == threesum_memoize(a, s, i+1, running, triplets, badstarts) - markbad(running) - end - end - if i == 1 - @show badstarts - return triplets - end + triples end -function threesum(a::AbstractVector{<:Integer}, s, i=1, running::Tuple=(), triplets=Set{Tuple}()) - if i > length(a) && length(running) < 3 - return - end - if length(running) == 3 - if sum(running) == s - push!(triplets, running) +function nsum(a::AbstractVector{T}, s, n) where {T<:Integer} + if n == 1 + tuple.(filter(==(s), a)) + elseif n == 2 + twosum_linear(a, s) + else + tuples = Tuple[] + for e = a + f = s - e + append!(tuples, ((e, p...) for p = nsum(a, f, n-1))) end - elseif length(running) == 2 - current = a[i] - # incl/excl/new - threesum(a, s, i+1, running, triplets) - threesum(a, s, i+1, (running..., current), triplets) - threesum(a, s, i+1, (running[1], current), triplets) - threesum(a, s, i+1, (running[2], current), triplets) - threesum(a, s, i+1, (), triplets) - elseif length(running) == 1 - current = a[i] - # incl/excl current element - threesum(a, s, i+1, (running[1], current), triplets) - threesum(a, s, i+1, running, triplets) - threesum(a, s, i+1, (), triplets) - elseif length(running) == 0 - current = a[i] - threesum(a, s, i+1, (current,), triplets) - threesum(a, s, i+1, running, triplets) - end - if i == 1 - return triplets + tuples end end