Python线性组合生成函数转Go函数：yield逻辑实现难题

原Python实现 #

def linear_compositions(n, k):
    if n < 0 or k < 0:
        return
    elif k == 0:  
        if n == 0:
            yield []
        return
    elif k == 1:
        yield [n]
        return
    else:
        for i in range(1, n):
            for comp in linear_compositions(n - i, k - 1):
                yield [i] + comp

问题描述 #

尝试将上述Python函数转换为Go函数时，编写的代码无法正确处理else块中的生成逻辑，输入n=12、k=5时返回大量索引0为8的切片集合，不符合预期。以下是错误的Go实现：

func linearComposition(n, k int) [][]int {
    var result [][]int
    if n < 0 || k < 0 {
        return nil
    } else if k == 0 {
        if n == 0 {
            return result
        }
    } else if k == 1 {
        return [][]int{{n}}
    } else {
        for i := 1; i < n; i++ {
            for _, comp := range linearComposition(n-1, k-1) {
                result = append(result, comp)
            }
        }
    }
    return result
}

错误分析与修正 #

原Go代码存在三个核心错误：

1. 递归参数错误：递归调用时传入了n-1，而非Python中的n-i，导致子组合的基数错误。
2. 组合拼接缺失：没有将当前循环的i添加到子组合的头部，对应Python代码中的[i] + comp逻辑。
3. 边界条件处理错误：当k=0且n=0时，原Python会生成空列表，而Go代码返回了空切片集合，应该返回包含空切片的集合[][]int{{}}。

修正后的Go实现 #

func linearComposition(n, k int) [][]int {
    var result [][]int
    if n < 0 || k < 0 {
        return nil
    }
    if k == 0 {
        if n == 0 {
            return [][]int{{}}
        }
        return result
    }
    if k == 1 {
        return [][]int{{n}}
    }
    // 处理k>1的情况
    for i := 1; i < n; i++ {
        // 递归获取n-i拆分为k-1份的组合
        subComps := linearComposition(n-i, k-1)
        for _, comp := range subComps {
            // 将当前i加到子组合头部，生成新组合
            newComp := append([]int{i}, comp...)
            result = append(result, newComp)
        }
    }
    return result
}

说明 #

• 递归调用时正确传入n-i，保证子组合的和为n-i，加上当前的i后总和为n。
• 使用append([]int{i}, comp...)完成Python中[i] + comp的拼接逻辑，生成新的组合切片。
• 修正k=0的边界条件，当n=0时返回包含空切片的集合，与原Python逻辑一致。

内容的提问来源于stack exchange，提问作者Enrico