Merge pull request #10 from randy3k/combinperm

jiahao · jiahao · commit ebf131911b04 · 2015-11-11T22:45:19.000-05:00
four more arrangement generators
diff --git a/src/combinations.jl b/src/combinations.jl
@@ -1,4 +1,7 @@
-export combinations, CoolLexCombinations
+export combinations,
+       CoolLexCombinations,
+       multiset_combinations,
+       with_replacement_combinations
 
 #The Combinations iterator
 
@@ -132,3 +135,118 @@ end
 done(C::CoolLexCombinations, S::CoolLexIterState) = (S.R3 & S.R2 != 0)
 
 length(C::CoolLexCombinations) = max(0, binomial(C.n, C.t))
+
+
+immutable MultiSetCombinations{T}
+    m::T
+    f::Vector{Int}
+    t::Int
+    ref::Vector{Int}
+end
+
+eltype{T}(::Type{MultiSetCombinations{T}}) = Vector{eltype(T)}
+
+function length(c::MultiSetCombinations)
+    t = c.t
+    if t > length(c.ref)
+        return 0
+    end
+    p = [1; zeros(Int,t)]
+    for i in 1:length(c.f)
+        f = c.f[i]
+        if i == 1
+            for j in 1:min(f, t)
+                p[j+1] = 1
+            end
+        else
+            for j in t:-1:1
+                p[j+1] = sum(p[max(1,j+1-f):(j+1)])
+            end
+        end
+    end
+    return p[t+1]
+end
+
+function multiset_combinations{T<:Integer}(m, f::Vector{T}, t::Integer)
+    length(m) == length(f) || error("Lengths of m and f are not the same.")
+    ref = length(f) > 0 ? vcat([[i for j in 1:f[i] ] for i in 1:length(f)]...) : Int[]
+    if t < 0
+        t = length(ref) + 1
+    end
+    MultiSetCombinations(m, f, t, ref)
+end
+
+"generate all combinations of size t from an array a with possibly duplicated elements."
+function multiset_combinations{T}(a::T, t::Integer)
+    m = unique(collect(a))
+    f = Int[sum([c == x for c in a]) for x in m]
+    multiset_combinations(m, f, t)
+end
+
+start(c::MultiSetCombinations) = c.ref
+function next(c::MultiSetCombinations, s)
+    ref = c.ref
+    n = length(ref)
+    t = c.t
+    changed = false
+    comb = [c.m[s[i]] for i in 1:t]
+    if t > 0
+        s = copy(s)
+        for i in t:-1:1
+            if s[i] < ref[i + (n - t)]
+                j = 1
+                while ref[j] <= s[i]; j += 1; end
+                s[i] = ref[j]
+                for l in (i+1):t
+                    s[l] = ref[j+=1]
+                end
+                changed = true
+                break
+            end
+        end
+        !changed && (s[1] = n+1)
+    else
+        s = [n+1]
+    end
+    (comb, s)
+end
+done(c::MultiSetCombinations, s) =
+    (!isempty(s) && max(s[1], c.t) > length(c.ref)) ||  (isempty(s) && c.t > 0)
+
+immutable WithReplacementCombinations{T}
+    a::T
+    t::Int
+end
+
+eltype{T}(::Type{WithReplacementCombinations{T}}) = Vector{eltype(T)}
+
+length(c::WithReplacementCombinations) = binomial(length(c.a)+c.t-1, c.t)
+
+"generate all combinations with replacement of size t from an array a."
+with_replacement_combinations(a, t::Integer) = WithReplacementCombinations(a, t)
+
+start(c::WithReplacementCombinations) = [1 for i in 1:c.t]
+function next(c::WithReplacementCombinations, s)
+    n = length(c.a)
+    t = c.t
+    comb = [c.a[si] for si in s]
+    if t > 0
+        s = copy(s)
+        changed = false
+        for i in t:-1:1
+            if s[i] < n
+                s[i] += 1
+                for j in (i+1):t
+                    s[j] = s[i]
+                end
+                changed = true
+                break
+            end
+        end
+        !changed && (s[1] = n+1)
+    else
+        s = [n+1]
+    end
+    (comb, s)
+end
+done(c::WithReplacementCombinations, s) = !isempty(s) && s[1] > length(c.a) || c.t < 0
diff --git a/src/permutations.jl b/src/permutations.jl
@@ -2,47 +2,133 @@
 
 export
     levicivita,
+    multiset_permutations,
     nthperm!,
     nthperm,
     parity,
     permutations
 
-#The basic permutations iterator
 
 immutable Permutations{T}
     a::T
+    t::Int
 end
 
+eltype{T}(::Type{Permutations{T}}) = Vector{eltype(T)}
+
+length(p::Permutations) = (0 <= p.t <= length(p.a))?factorial(length(p.a), length(p.a)-p.t):0
+
 """
 Generate all permutations of an indexable object. Because the number of permutations can be very large, this function returns an iterator object. Use `collect(permutations(array))` to get an array of all permutations.
 """
-permutations(a) = Permutations(a)
-
-eltype{T}(::Type{Permutations{T}}) = Vector{eltype(T)}
+permutations(a) = Permutations(a, length(a))
 
-length(p::Permutations) = factorial(length(p.a))
+"""
+Generate all size t permutations of an indexable object.
+"""
+function permutations(a, t::Integer)
+    if t < 0
+        t = length(a) + 1
+    end
+    Permutations(a, t)
+end
 
 start(p::Permutations) = [1:length(p.a);]
-function next(p::Permutations, s)
-    perm = [p.a[si] for si in s]
-    if isempty(p.a)
-        # special case to generate 1 result for len==0
-        return (perm,[1])
+next(p::Permutations, s) = nextpermutation(p.a, p.t ,s)
+
+function nextpermutation(m, t, s)
+    perm = [m[s[i]] for i in 1:t]
+    n = length(s)
+    if t <= 0
+        return(perm, [n+1])
     end
     s = copy(s)
-    k = length(s)-1
-    while k > 0 && s[k] > s[k+1];  k -= 1;  end
-    if k == 0
-        s[1] = length(s)+1   # done
+    if t < n
+        j = t + 1
+        while j <= n &&  s[t] >= s[j]; j+=1; end
+    end
+    if t < n && j <= n
+        s[t], s[j] = s[j], s[t]
+    else
+        if t < n
+            reverse!(s, t+1)
+        end
+        i = t - 1
+        while i>=1 && s[i] >= s[i+1]; i -= 1; end
+        if i > 0
+            j = n
+            while j>i && s[i] >= s[j]; j -= 1; end
+            s[i], s[j] = s[j], s[i]
+            reverse!(s, i+1)
+        else
+            s[1] = n+1
+        end
+    end
+    return(perm, s)
+end
+
+done(p::Permutations, s) = !isempty(s) && max(s[1], p.t) > length(p.a) ||  (isempty(s) && p.t > 0)
+
+immutable MultiSetPermutations{T}
+    m::T
+    f::Vector{Int}
+    t::Int
+    ref::Vector{Int}
+end
+
+eltype{T}(::Type{MultiSetPermutations{T}}) = Vector{eltype(T)}
+
+function length(c::MultiSetPermutations)
+    t = c.t
+    if t > length(c.ref)
+        return 0
+    end
+    if t > 20
+        g = [factorial(big(i)) for i in 0:t]
     else
-        l = length(s)
-        while s[k] >= s[l];  l -= 1;  end
-        s[k],s[l] = s[l],s[k]
-        reverse!(s,k+1)
+        g = [factorial(i) for i in 0:t]
+    end
+    p = [g[t+1]; zeros(Float64,t)]
+    for i in 1:length(c.f)
+        f = c.f[i]
+        if i == 1
+            for j in 1:min(f, t)
+                p[j+1] = g[t+1]/g[j+1]
+            end
+        else
+            for j in t:-1:1
+                q = 0
+                for k in (j+1):-1:max(1,j+1-f)
+                    q += p[k]/g[j+2-k]
+                end
+                p[j+1] = q
+            end
+        end
     end
-    (perm,s)
+    return round(Int, p[t+1])
 end
-done(p::Permutations, s) = !isempty(s) && s[1] > length(p.a)
+
+"generate all permutations of size t from an array a with possibly duplicated elements."
+function multiset_permutations{T<:Integer}(m, f::Vector{T}, t::Integer)
+    length(m) == length(f) || error("Lengths of m and f are not the same.")
+    ref = length(f) > 0 ? vcat([[i for j in 1:f[i] ] for i in 1:length(f)]...) : Int[]
+    if t < 0
+        t = length(ref) + 1
+    end
+    MultiSetPermutations(m, f, t, ref)
+end
+
+function multiset_permutations{T}(a::T, t::Integer)
+    m = unique(collect(a))
+    f = [sum([c == x for c in a]) for x in m]
+    multiset_permutations(m, f, t)
+end
+
+start(p::MultiSetPermutations) = p.ref
+next(p::MultiSetPermutations, s) = nextpermutation(p.m, p.t, s)
+done(p::MultiSetPermutations, s) =
+    !isempty(s) && max(s[1], p.t) > length(p.ref) ||  (isempty(s) && p.t > 0)
+
 
 
 "In-place version of nthperm."
@@ -136,4 +222,3 @@ function parity{T<:Integer}(p::AbstractVector{T})
     epsilon == 0 && throw(ArgumentError("Not a permutation"))
     epsilon == 1 ? 0 : 1
 end
-
diff --git a/test/combinations.jl b/test/combinations.jl
@@ -15,9 +15,29 @@ import Combinatorics: combinations
 
 @test collect(filter(x->(iseven(x[1])),combinations([1,2,3],2))) == Any[[2,3]]
 
+# multiset_combinations
+@test collect(multiset_combinations("aabc", 5)) == Any[]
+@test collect(multiset_combinations("aabc", 2)) == Any[['a','a'],['a','b'],['a','c'],['b','c']]
+@test collect(multiset_combinations("aabc", 1)) == Any[['a'],['b'],['c']]
+@test collect(multiset_combinations("aabc", 0)) == Any[Char[]]
+@test collect(multiset_combinations("aabc", -1)) == Any[]
+@test collect(multiset_combinations("", 1)) == Any[]
+@test collect(multiset_combinations("", 0)) == Any[Char[]]
+@test collect(multiset_combinations("", -1)) == Any[]
+
+# with_replacement_combinations
+@test collect(with_replacement_combinations("abc", 2)) == Any[['a','a'],['a','b'],['a','c'],
+                                                              ['b','b'],['b','c'],['c','c']]
+@test collect(with_replacement_combinations("abc", 1)) == Any[['a'],['b'],['c']]
+@test collect(with_replacement_combinations("abc", 0)) == Any[Char[]]
+@test collect(with_replacement_combinations("abc", -1)) == Any[]
+@test collect(with_replacement_combinations("", 1)) == Any[]
+@test collect(with_replacement_combinations("", 0)) == Any[Char[]]
+@test collect(with_replacement_combinations("", -1)) == Any[]
+
+
 #cool-lex iterator
 @test_throws DomainError collect(CoolLexCombinations(-1, 1))
 @test_throws DomainError collect(CoolLexCombinations(5, 0))
 @test collect(CoolLexCombinations(4,2)) == Vector[[1,2], [2,3], [1,3], [2,4], [3,4], [1,4]]
 @test isa(start(CoolLexCombinations(1000, 20)), Combinatorics.CoolLexIterState{BigInt})
-
diff --git a/test/permutations.jl b/test/permutations.jl
@@ -3,6 +3,7 @@ using Base.Test
 
 import Combinatorics: levicivita, nthperm, nthperm!, parity, permutations
 
+# permutations
 @test collect(permutations("abc")) == Any[['a','b','c'],['a','c','b'],['b','a','c'],
                                           ['b','c','a'],['c','a','b'],['c','b','a']]
 
@@ -11,6 +12,26 @@ import Combinatorics: levicivita, nthperm, nthperm!, parity, permutations
 
 @test length(permutations(0)) == 1
 
+@test collect(permutations("abc", 4)) == Any[]
+@test collect(permutations("abc", 2)) == Any[['a','b'],['a','c'],['b','a'],
+                                             ['b','c'],['c','a'],['c','b']]
+@test collect(permutations("abc", 0)) == Any[Char[]]
+@test collect(permutations("abc", -1)) == Any[]
+@test collect(permutations("", 1)) == Any[]
+@test collect(permutations("", 0)) == Any[Char[]]
+@test collect(permutations("", -1)) == Any[]
+
+# multiset_permutations
+@test collect(multiset_permutations("aabc", 5)) == Any[]
+@test collect(multiset_permutations("aabc", 2)) == Any[['a','a'],['a','b'], ['a','c'],['b','a'],
+                                                       ['b','c'],['c','a'],['c','b']]
+@test collect(multiset_permutations("aabc", 0)) == Any[Char[]]
+@test collect(multiset_permutations("aabc", -1)) == Any[]
+@test collect(multiset_permutations("", 1)) == Any[]
+@test collect(multiset_permutations("", 0)) == Any[Char[]]
+@test collect(multiset_permutations("", -1)) == Any[]
+@test length(multiset_permutations("aaaaaaaaaaaaaaaaaaaaab", 21)) == 22
+
 #nthperm!
 for n = 0:7, k = 1:factorial(n)
     p = nthperm!([1:n;], k)
@@ -32,4 +53,3 @@ end
 
 @test Combinatorics.nsetpartitions(-1) == 0
 @test collect(permutations([])) == Vector[[]]
-
