JuliaMath / Combinatorics.jl

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

with_replacement_permutations #118

Open astrozot opened 2 years ago

astrozot commented 2 years ago

I believe it would be useful to add a method to generate permutations with repetitions, so that, for example

julia> collect(with_replacement_permutations(1:2, 2))
4-element Vector{Vector{Int64}}:
 [1, 1]
 [1, 2]
 [2, 1]
 [2, 2]

A possible implementation could be the following

struct WithReplacementPermutations{T}
    a::T
    t::Int
end

Base.eltype(::Type{WithReplacementPermutations{T}}) where {T} = Vector{eltype(T)}

Base.length(p::WithReplacementPermutations) = length(p.a)^p.t

"""
    with_replacement_permutations(a, t)

Generate all permutations with replacement of size `t` from an array `a`.
"""
with_replacement_permutations(elements::T, len::Integer) where {T} =
    WithReplacementPermutations{T}(unique(elements), len)

function Base.iterate(p::WithReplacementPermutations, s = ones(Int, p.t))
    (!isempty(s) && s[end] > length(p.a) || (isempty(s) && p.t > 0)) && return
    nextpermutationwithreplacement(p.a, p.t, s)
end

function nextpermutationwithreplacement(a, t, state)
    cur = a[state]
    n = length(state)
    state[1] += 1
    local i = 1
    while i < n && state[i] > length(a)
        state[i] = 1
        state[i+1] += 1
        i += 1
    end
    return cur, state
end