Skip to content

Suggestion: combinations as bitmasks #163

Open
@sgaure

Description

@sgaure

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions