mirror of
https://github.com/gonum/gonum.git
synced 2025-10-06 15:47:01 +08:00
combin: Add helpers for dealing with permutations (#1076)
* combin: Add helpers for dealing with permutations
This commit is contained in:
Brendan Tracey
@@ -72,8 +72,8 @@ func LogGeneralizedBinomial(n, k float64) float64 {
|
||||
return a - b - c
|
||||
}
|
||||
|
||||
// CombinationGenerator generates combinations iteratively. Combinations may be
|
||||
// called to generate all combinations collectively.
|
||||
// CombinationGenerator generates combinations iteratively. The Combinations
|
||||
// function may be called to generate all combinations collectively.
|
||||
type CombinationGenerator struct {
|
||||
n int
|
||||
k int
|
||||
@@ -118,29 +118,28 @@ func (c *CombinationGenerator) Next() bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// Combination generates the current combination. If next is non-nil, it must have
|
||||
// length k and the result will be stored in-place into combination. If combination
|
||||
// Combination returns the current combination. If dst is non-nil, it must have
|
||||
// length k and the result will be stored in-place into dst. If dst
|
||||
// is nil a new slice will be allocated and returned. If all of the combinations
|
||||
// have already been constructed (Next() returns false), Combination will panic.
|
||||
//
|
||||
// Next must be called to initialize the first value before calling Combination
|
||||
// or Combination will panic. The value returned by Combination is only changed
|
||||
// during calls to Next.
|
||||
func (c *CombinationGenerator) Combination(combination []int) []int {
|
||||
func (c *CombinationGenerator) Combination(dst []int) []int {
|
||||
if c.remaining == -1 {
|
||||
panic("combin: all combinations have been generated")
|
||||
}
|
||||
if c.previous == nil {
|
||||
panic("combin: Combination called before Next")
|
||||
}
|
||||
if combination == nil {
|
||||
combination = make([]int, c.k)
|
||||
}
|
||||
if len(combination) != c.k {
|
||||
if dst == nil {
|
||||
dst = make([]int, c.k)
|
||||
} else if len(dst) != c.k {
|
||||
panic(badInput)
|
||||
}
|
||||
copy(combination, c.previous)
|
||||
return combination
|
||||
copy(dst, c.previous)
|
||||
return dst
|
||||
}
|
||||
|
||||
// Combinations generates all of the combinations of k elements from a
|
||||
@@ -150,7 +149,7 @@ func (c *CombinationGenerator) Combination(combination []int) []int {
|
||||
// n and k must be non-negative with n >= k, otherwise Combinations will panic.
|
||||
//
|
||||
// CombinationGenerator may alternatively be used to generate the combinations
|
||||
// iteratively instead of collectively.
|
||||
// iteratively instead of collectively, or IndexToCombination for random access.
|
||||
func Combinations(n, k int) [][]int {
|
||||
combins := Binomial(n, k)
|
||||
data := make([][]int, combins)
|
||||
@@ -368,3 +367,267 @@ func SubFor(sub []int, idx int, dims []int) []int {
|
||||
sub[len(sub)-1] = idx
|
||||
return sub
|
||||
}
|
||||
|
||||
// NumPermutations returns the number of permutations when selecting k
|
||||
// objects from a set of n objects when the selection order matters.
|
||||
// No check is made for overflow.
|
||||
//
|
||||
// NumPermutations panics if either n or k is negative, or if k is
|
||||
// greater than n.
|
||||
func NumPermutations(n, k int) int {
|
||||
if n < 0 {
|
||||
panic("combin: n is negative")
|
||||
}
|
||||
if k < 0 {
|
||||
panic("combin: k is negative")
|
||||
}
|
||||
if k > n {
|
||||
panic("combin: k is greater than n")
|
||||
}
|
||||
p := 1
|
||||
for i := n - k + 1; i <= n; i++ {
|
||||
p *= i
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// Permutations generates all of the permutations of k elements from a
|
||||
// set of size n. The returned slice has length NumPermutations(n, k)
|
||||
// and each inner slice has length k.
|
||||
//
|
||||
// n and k must be non-negative with n >= k, otherwise Permutations will panic.
|
||||
//
|
||||
// PermutationGenerator may alternatively be used to generate the permutations
|
||||
// iteratively instead of collectively, or IndexToPermutation for random access.
|
||||
func Permutations(n, k int) [][]int {
|
||||
nPerms := NumPermutations(n, k)
|
||||
data := make([][]int, nPerms)
|
||||
if len(data) == 0 {
|
||||
return data
|
||||
}
|
||||
for i := 0; i < nPerms; i++ {
|
||||
data[i] = IndexToPermutation(nil, i, n, k)
|
||||
}
|
||||
return data
|
||||
}
|
||||
|
||||
// PermutationGenerator generates permutations iteratively. The Permutations
|
||||
// function may be called to generate all permutations collectively.
|
||||
type PermutationGenerator struct {
|
||||
n int
|
||||
k int
|
||||
nPerm int
|
||||
idx int
|
||||
permutation []int
|
||||
}
|
||||
|
||||
// NewPermutationGenerator returns a PermutationGenerator for generating the
|
||||
// permutations of k elements from a set of size n.
|
||||
//
|
||||
// n and k must be non-negative with n >= k, otherwise NewPermutationGenerator
|
||||
// will panic.
|
||||
func NewPermutationGenerator(n, k int) *PermutationGenerator {
|
||||
return &PermutationGenerator{
|
||||
n: n,
|
||||
k: k,
|
||||
nPerm: NumPermutations(n, k),
|
||||
idx: -1,
|
||||
permutation: make([]int, k),
|
||||
}
|
||||
}
|
||||
|
||||
// Next advances the iterator if there are permutations remaining to be generated,
|
||||
// and returns false if all permutations have been generated. Next must be called
|
||||
// to initialize the first value before calling Permutation or Permutation will
|
||||
// panic. The value returned by Permutation is only changed during calls to Next.
|
||||
func (p *PermutationGenerator) Next() bool {
|
||||
if p.idx >= p.nPerm-1 {
|
||||
p.idx = p.nPerm // so Permutation can panic.
|
||||
return false
|
||||
}
|
||||
p.idx++
|
||||
IndexToPermutation(p.permutation, p.idx, p.n, p.k)
|
||||
return true
|
||||
}
|
||||
|
||||
// Permutation returns the current permutation. If dst is non-nil, it must have
|
||||
// length k and the result will be stored in-place into dst. If dst
|
||||
// is nil a new slice will be allocated and returned. If all of the permutations
|
||||
// have already been constructed (Next() returns false), Permutation will panic.
|
||||
//
|
||||
// Next must be called to initialize the first value before calling Permutation
|
||||
// or Permutation will panic. The value returned by Permutation is only changed
|
||||
// during calls to Next.
|
||||
func (p *PermutationGenerator) Permutation(dst []int) []int {
|
||||
if p.idx == p.nPerm {
|
||||
panic("combin: all permutations have been generated")
|
||||
}
|
||||
if p.idx == -1 {
|
||||
panic("combin: Permutation called before Next")
|
||||
}
|
||||
if dst == nil {
|
||||
dst = make([]int, p.k)
|
||||
} else if len(dst) != p.k {
|
||||
panic(badInput)
|
||||
}
|
||||
copy(dst, p.permutation)
|
||||
return dst
|
||||
}
|
||||
|
||||
// PermutationIndex returns the index of the given permutation.
|
||||
//
|
||||
// The functions PermutationIndex and IndexToPermutation define a bijection
|
||||
// between the integers and the NumPermutations(n, k) number of possible permutations.
|
||||
// PermutationIndex returns the inverse of IndexToPermutation.
|
||||
//
|
||||
// PermutationIndex panics if perm is not a permutation of k of the first
|
||||
// [0,n) integers, if n or k are non-negative, or if k is greater than n.
|
||||
func PermutationIndex(perm []int, n, k int) int {
|
||||
if n < 0 || k < 0 {
|
||||
panic(badNegInput)
|
||||
}
|
||||
if n < k {
|
||||
panic(badSetSize)
|
||||
}
|
||||
if len(perm) != k {
|
||||
panic("combin: bad length permutation")
|
||||
}
|
||||
contains := make(map[int]struct{}, k)
|
||||
for _, v := range perm {
|
||||
if v < 0 || v >= n {
|
||||
panic("combin: bad element")
|
||||
}
|
||||
contains[v] = struct{}{}
|
||||
}
|
||||
if len(contains) != k {
|
||||
panic("combin: perm contains non-unique elements")
|
||||
}
|
||||
if n == k {
|
||||
// The permutation is the ordering of the elements.
|
||||
return equalPermutationIndex(perm)
|
||||
}
|
||||
|
||||
// The permutation index is found by finding the combination index and the
|
||||
// equalPermutation index. The combination index is found by just sorting
|
||||
// the elements, and the the permutation index is the ordering of the size
|
||||
// of the elements.
|
||||
tmp := make([]int, len(perm))
|
||||
copy(tmp, perm)
|
||||
idx := make([]int, len(perm))
|
||||
for i := range idx {
|
||||
idx[i] = i
|
||||
}
|
||||
s := sortInts{tmp, idx}
|
||||
sort.Sort(s)
|
||||
order := make([]int, len(perm))
|
||||
for i, v := range idx {
|
||||
order[v] = i
|
||||
}
|
||||
combIdx := CombinationIndex(tmp, n, k)
|
||||
permIdx := equalPermutationIndex(order)
|
||||
return combIdx*NumPermutations(k, k) + permIdx
|
||||
}
|
||||
|
||||
type sortInts struct {
|
||||
data []int
|
||||
idx []int
|
||||
}
|
||||
|
||||
func (s sortInts) Len() int {
|
||||
return len(s.data)
|
||||
}
|
||||
|
||||
func (s sortInts) Less(i, j int) bool {
|
||||
return s.data[i] < s.data[j]
|
||||
}
|
||||
|
||||
func (s sortInts) Swap(i, j int) {
|
||||
s.data[i], s.data[j] = s.data[j], s.data[i]
|
||||
s.idx[i], s.idx[j] = s.idx[j], s.idx[i]
|
||||
}
|
||||
|
||||
// IndexToPermutation returns the permutation corresponding to the given index.
|
||||
//
|
||||
// The functions PermutationIndex and IndexToPermutation define a bijection
|
||||
// between the integers and the NumPermutations(n, k) number of possible permutations.
|
||||
// IndexToPermutation returns the inverse of PermutationIndex.
|
||||
//
|
||||
// The permutation is stored in-place into dst if dst is non-nil, otherwise
|
||||
// a new slice is allocated and returned.
|
||||
//
|
||||
// IndexToPermutation panics if n or k are non-negative, if k is greater than n,
|
||||
// or if idx is not in [0, NumPermutations(n,k)-1]. IndexToPermutation will also panic
|
||||
// if dst is non-nil and len(dst) is not k.
|
||||
func IndexToPermutation(dst []int, idx, n, k int) []int {
|
||||
nPerm := NumPermutations(n, k)
|
||||
if idx < 0 || idx >= nPerm {
|
||||
panic("combin: invalid index")
|
||||
}
|
||||
if dst == nil {
|
||||
dst = make([]int, k)
|
||||
} else if len(dst) != k {
|
||||
panic(badInput)
|
||||
}
|
||||
if n == k {
|
||||
indexToEqualPermutation(dst, idx)
|
||||
return dst
|
||||
}
|
||||
|
||||
// First, we index into the combination (which of the k items to choose)
|
||||
// and then we index into the n == k permutation of those k items. The
|
||||
// indexing acts like a matrix with nComb rows and factorial(k) columns.
|
||||
kPerm := NumPermutations(k, k)
|
||||
combIdx := idx / kPerm
|
||||
permIdx := idx % kPerm
|
||||
comb := IndexToCombination(nil, combIdx, n, k) // Gives us the set of integers.
|
||||
perm := make([]int, len(dst))
|
||||
indexToEqualPermutation(perm, permIdx) // Gives their order.
|
||||
for i, v := range perm {
|
||||
dst[i] = comb[v]
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
// equalPermutationIndex returns the index of the given permutation of the
|
||||
// first k integers.
|
||||
func equalPermutationIndex(perm []int) int {
|
||||
// Note(btracey): This is an n^2 algorithm, but factorial increases
|
||||
// very quickly (25! overflows int64) so this is not a problem in
|
||||
// practice.
|
||||
idx := 0
|
||||
for i, u := range perm {
|
||||
less := 0
|
||||
for _, v := range perm[i:] {
|
||||
if v < u {
|
||||
less++
|
||||
}
|
||||
}
|
||||
idx += less * factorial(len(perm)-i-1)
|
||||
}
|
||||
return idx
|
||||
}
|
||||
|
||||
// indexToEqualPermutation returns the permutation for the first len(dst)
|
||||
// integers for the given index.
|
||||
func indexToEqualPermutation(dst []int, idx int) {
|
||||
for i := range dst {
|
||||
dst[i] = i
|
||||
}
|
||||
for i := range dst {
|
||||
f := factorial(len(dst) - i - 1)
|
||||
r := idx / f
|
||||
v := dst[i+r]
|
||||
copy(dst[i+1:i+r+1], dst[i:i+r])
|
||||
dst[i] = v
|
||||
idx %= f
|
||||
}
|
||||
}
|
||||
|
||||
// factorial returns a!.
|
||||
func factorial(a int) int {
|
||||
f := 1
|
||||
for i := 2; i <= a; i++ {
|
||||
f *= i
|
||||
}
|
||||
return f
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user