sig-pub/pkg/utils/collect/combine.go

package collect

// CartesianCount 笛卡尔积组合数
func CartesianCount[T any](sets ...[]T) (total int) {
	total = 1
	for _, set := range sets {
		total *= len(set)
	}
	return
}

// CartesianProduct 笛卡尔积(返回所有可能的组合)
// 输入: [][]T 类型的二维切片
// 输出: []T 类型的切片组成的切片
func CartesianProduct[T any](sets ...[]T) [][]T {
	if len(sets) == 0 {
		return [][]T{{}}
	}

	// 初始化为第一个集合的所有单元素组合
	result := make([][]T, len(sets[0]))
	for i, v := range sets[0] {
		result[i] = []T{v}
	}

	// 依次处理后续每一组
	for _, set := range sets[1:] {
		if len(set) == 0 {
			return [][]T{}
		}

		newResult := make([][]T, 0, len(result)*len(set))

		for _, prev := range result {
			for _, curr := range set {
				// 预分配空间，避免频繁扩容
				combined := make([]T, 0, len(prev)+1)
				combined = append(combined, prev...)
				combined = append(combined, curr)
				newResult = append(newResult, combined)
			}
		}

		result = newResult
	}

	return result
}

// CartesianYield 生成笛卡尔积的所有组合，并为每个组合调用回调函数
// callback: func(comb []T) bool - 处理当前组合，返回 true 继续生成，返回 false 停止
// return: 生成的组合总数
func CartesianYield[T any](sets [][]T, callback func([]T) bool) int {
	if len(sets) == 0 {
		return 0
	}

	// 检查是否有空集
	for _, s := range sets {
		if len(s) == 0 {
			return 0
		}
	}

	// 初始化索引计数器
	indices := make([]int, len(sets))
	comb := make([]T, len(sets)) // 复用缓冲区，避免每次分配

	count := 0
	for {
		// 构建当前组合
		for i, idx := range indices {
			comb[i] = sets[i][idx]
		}

		// 调用回调
		count++
		if !callback(comb) {
			break // 停止生成
		}

		// 递增索引，像计数器一样
		i := len(indices) - 1
		for i >= 0 {
			indices[i]++
			if indices[i] < len(sets[i]) {
				break
			}
			indices[i] = 0
			i--
		}

		if i < 0 {
			break // 所有组合已生成
		}
	}

	return count
}