Zero allocated iterator for powersets

Starting with the empty set, Gray codes can be used to determine the next element to add or remove so that you eventually hit every subset exactly once. Unfortunately, I’m having trouble creating an iterator for this setup because I want it to output the subsets each iteration. My current implementation has to call an update function inside of the loop.

#=
Iterating through a powerset
Idea: Start with an empty subset and use gray codes to determine the next element to add or remove 

Pros: 
    - no allocations
    - Can be 8x faster than Combinatorics.jl
Cons:
    - subsets are not ordered by size (monotonic gray codes?)
    - output type is a Set (meaning no powersets of repeated values. e.g., [1,1,2])
    - may be impossible to create an iterator with no allocations
=#

#Which bit to flip when creating the binary reflected Gray code
#Add 1 for base-1 indexing
bit2flip(n::T) where T<:Integer = trailing_zeros(n) + one(T)

#Structure to hold current subset and count how many subsets have been calculated
mutable struct powerSetIter{V<:AbstractVector, S<:AbstractSet}
    set::V
    setLength::Int
    subset::S
    count::Int
end

#Give a size hint for the subset
function powerSetIter(set::AbstractVector{T}) where T
    n = length(set)
    subset = Set{T}()
    sizehint!(subset,n) 
    
    return powerSetIter(set,2^n,subset,1)
end

#Helpful if you want a powerset of something like 1:n
function powerSetIter(set::AbstractVector{T}) where T<:Integer
    n = length(set)
    subset = BitSet()
    sizehint!(subset,n) #might be unnecessary
    
    return powerSetIter(set,2^n,subset,1)
end

#Given the current state calculate the next subset
#Can this be turned into an iterator?????
function nextSubset!(S::powerSetIter)
    
    #loop back to start
    if S.count == S.setLength
        S.count = 1
        empty!(S.subset)
        return nothing
    end
    
    element = S.set[bit2flip(S.count)]
    S.count += 1
    
    #Add or remove element (fast for Sets)
    if element ∈ S.subset
        delete!(S.subset,element)
    else
        push!(S.subset,element)
    end
    
    return nothing
end

A for-loop would look something like this:

S = powerSetIter([:a,:b,:c])

for _ in 1:S.setLength
    println(S.subset)
    nextSubset!(S)
end

Set{Symbol}()
Set([:a])
Set([:a, :b])
Set([:b])
Set([:b, :c])
Set([:a, :b, :c])
Set([:a, :c])
Set([:c])

And as mentioned there are no allocations:

using BenchmarkTools

#Simple function to iterate through powersets
function setIterate(S::powerSetIter)
    l = 0
    for _ in 1:S.setLength
        l += length(S.subset)
        nextSubset!(S)
    end

    return l
end

s = powerSetIter(1:20)

@benchmark setIterate($s)

BenchmarkTools.Trial: 1004 samples with 1 evaluation.
 Range (min â€Ķ max):  4.509 ms â€Ķ   8.825 ms  ┊ GC (min â€Ķ max): 0.00% â€Ķ 0.00%
 Time  (median):     4.915 ms               ┊ GC (median):    0.00%
 Time  (mean Âą σ):   4.975 ms Âą 376.404 Ξs  ┊ GC (mean Âą σ):  0.00% Âą 0.00%

   ▄▂▂██▆▄▃▅▇▂▅▅▅▂▅▂▁ ▃
  ▇██████████████████▇█████▆▅▄▃▄▃▂▄▃▃▃▃▃▃▂▁▁▃▁▁▃▂▃▁▂▁▂▁▁▃▁▁▁▃ ▅
  4.51 ms         Histogram: frequency by time        6.44 ms <

 Memory estimate: 0 bytes, allocs estimate: 0.

Finally here is the comparision to Combinatorics.jl

1 Like