JuliaMath / Combinatorics.jl

A combinatorics library for Julia
http://juliamath.github.io/Combinatorics.jl/dev/
Other
214 stars 58 forks source link

Suggestion: combinations as bitmasks #163

Open sgaure opened 4 months ago

sgaure commented 4 months ago

I sometimes need all subsets of a particular size of some small integers. Moreover, I need them as bit masks, all bit patterns of a particular weight, e.g. 0b011, 0b101, 0b110. And it typically needs to run fast without allocations. I've made an iterator for this purpose. I came to think that it might fit well into the Combinatorics package. Here it is, if there's interest for incorporating it. A similar function for generating all subsets (i.e. all integers of all weights) shouldn't be needed, it's just counting.

struct BitCombinations{T, T1, T2}
    weight::T1
    width::T2
    thelast::T
    function BitCombinations(weight, width, ::Type{T}) where {T}
        isbitstype(T) || throw(ArgumentError("$T is not a bitstype"))
        T1 = typeof(weight)
        T2 = typeof(width)
        new{T,T1,T2}(weight, width, ((-1 % T) >>> (8*sizeof(T) - weight)) << (width-weight))
    end
end
Base.length(fw::BitCombinations) = binomial(fw.width, fw.weight)
Base.IteratorEltype(::Type{<:BitCombinations}) = Base.HasEltype()
Base.eltype(::Type{<:BitCombinations{T}}) where T = T

# from https://graphics.stanford.edu/~seander/bithacks.html#NextBitPermutation
@inline function next_combination(x::T) where T
    t = x | (x-one(T))
    return (t+one(T)) | (((~t & -~t) - one(T)) >>> (trailing_zeros(x) + one(T)))
end

@inline function Base.iterate(fw::BitCombinations{T}) where T
    val = (-1 % T) >>> (8*sizeof(T) - fw.weight)
    return (val,val)
end

@inline function Base.iterate(fw::BitCombinations, state)
    state == fw.thelast && return nothing
    val = next_combination(state)
    return (val,val)
end

"""
bitcombinations(weight, width, ::Type{T}=UInt64)

Generate all bit patterns with `weight` bits set in a bit field of width `width`.
An iterator which returns integers of type `T` is created. The type `T` must be a bits type.
The patterns are returned in lexicographical order, that is, in arithmetic order.

Such combinations can also be generated by [`combinations`](@ref), but in case the
actual bit patterns are needed, this function is faster and non-allocating.
"""
function bitcombinations(weight::Integer, width::Integer, ::Type{T}=UInt64) where {T<:Integer}
    nbits = 8*sizeof(T)
    (0 ≤ weight ≤ width ≤ nbits) || 
        throw(ArgumentError(lazy"0 ≤ weight($weight) ≤ width($width) ≤ 8*sizeof($T)($nbits) failed"))
    BitCombinations(weight, width, T)
end