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