golang实现 python itertools的组合

之前写过一个版本,代码太多了。今天参照 python的itertools.combinations源码重新整理了一版。

python

def combinations(iterable, r):
    # combinations('ABCD', 2) --> AB AC AD BC BD CD
    # combinations(range(4), 3) --> 012 013 023 123
    pool = tuple(iterable)
    n = len(pool)
    if r > n:
        return
    indices = list(range(r))
    print(indices, range(r))
    yield tuple(pool[i] for i in indices)
    while True:
        re = list(reversed(range(r)))
        # print(re)
        for i in re:
            # print(i, indices[i])
            if indices[i] != i + n - r:
                break
        else:
            return
        indices[i] += 1
        for j in range(i + 1, r):
            indices[j] = indices[j - 1] + 1
        yield tuple(pool[i] for i in indices)

golang

package main

import (
	"fmt"
	"math/big"
	"sort"
	"strconv"

	"gonum.org/v1/gonum/mat"
	"gonum.org/v1/gonum/stat"
)

// 组合
func increCombine(nums []interface{}, r int) (result [][]interface{}) {
	n := len(nums)
	if r > n {
		return
	}
	fr := increCobineAmount(n, r)
	indices := make([]int, r)
	indicesReverse := make([]int, r)
	for i := 0; i < r; i++ {
		indices[i] = i
		indicesReverse[i] = i
	}
	// 获取第一个
	first := nums[indices[0]:r]
	result = append(result, first)
	if r == n {
		return
	}
	sort.Sort(sort.Reverse(sort.IntSlice(indicesReverse)))
	lr := len(result)
	for lr < int(fr.Int64()) {
		i := 0
		for _, ii := range indicesReverse {

			i = ii
			if indices[i] != i+n-r {

				break
			}
		}
		indices[i] += 1
		var item []interface{}
		for j := i + 1; j < r; j++ {
			indices[j] = indices[j-1] + 1
		}
		for _, ti := range indices {
			item = append(item, nums[ti])
		}
		result = append(result, item)
		lr = len(result)
	}
	return
}

// 组合量n!/r!/(n-r)!
func increCobineAmount(total int, gNum int) *big.Int {
	cNums := factorial(big.NewInt(int64(total)))
	cNums = cNums.Div(cNums, factorial(big.NewInt(int64(gNum))))
	cNums = cNums.Div(cNums, factorial(big.NewInt(int64(total-gNum))))
	return cNums

}

// 阶乘n!
func factorial(n *big.Int) *big.Int {

	if n.Int64() == 1 {
		return big.NewInt(1)
	}
	return n.Mul(n, factorial(big.NewInt(n.Int64()-1)))
}

func main() {
	f := factorial(big.NewInt(10))
	f1 := increCobineAmount(33, 6)
	fmt.Println(f, f1.Int64())

	var a []interface{}
	for i := 1; i < 10; i++ {
		a = append(a, strconv.Itoa(i)+"aa")
	}
	// fmt.Println(a)
	// as = []interface{1, 2, 3, 4}
	resp := increCombine(a, 3)
	fmt.Println("resu:", len(resp))
}

结果

[[1aa 2aa 3aa] [1aa 2aa 4aa] [1aa 2aa 5aa] [1aa 3aa 4aa] [1aa 3aa 5aa] [1aa 4aa 5aa] [2aa 3aa 4aa] [2aa 3aa 5aa] [2aa 4aa 5aa] [3aa 4aa 5aa]]

性能比上一版本提示很多,

#C10,3 
	resu: 84
	cost: 72.583µs
#C100,3
	resu: 156849
	cost: 53.556291ms
#C1000,3
	resu: 165668499
	cost: 1m45.838739166s

再说一下为什么用math/big来计算阶乘:

int64会溢出,输出的数字不准确

发表评论
留言与评论(共有 0 条评论) “”
   
验证码:

相关文章

推荐文章