n选k组合问题——python算法

75 阅读1分钟

引入

面试算法题中遇到一个需要遍历n选k所有组合的问题,在此分享一下解决组合问题的算法。

组合问题

组合(combination)是一个数学名词。一般地,从n个不同的元素中,任取m(m≤n)个元素为一组,叫作从n个不同元素中取出m个元素的一个组合。我们把有关求组合的个数的问题叫作组合问题。

高中学过排列组合,组合数公式:Cnk=n(n1)(n2)...(nk+1)k!C_n^k=\frac{n(n-1)(n-2)...(n-k+1)}{k!}

算法思路

理解组合问题不难,但是算法思路可能一时间想不到。

这个问题可以理解为子集枚举,假设我们需要找到一个长度为n的序列a的所有子序列

代码框架如下:

def dfs(cur, n) {
    if (cur == n + 1) {
        # 记录答案
        res.append(tmp)
        return
    }
    # 选择当前位置
    temp.append(cur);
    dfs(cur + 1, n);
    
    # 不选择当前位置
    temp.pop();
    dfs(cur + 1, n);
}

代码的关键点在于对下一位,选择当前位置不选择当前位置

此时组合问题只是多了一个条件,子序列的长度要为k,将其加入终止条件即可

代码实现

# 返回所有0-n-1所有可能的下标排列
def combine(n, k):
    res, tmp = [], []
    if n == k: return [list(range(n))]
    if n < k: return []

    def dfs(i):  # 对应下标为i的数
        if len(tmp) + (n - i) < k:  # 剪枝:无法构建长度为k的组合
            return
        if len(tmp) == k:
            res.append(tmp[:])
            return
            # 加入下标为k的数
        tmp.append(i)
        dfs(i + 1)
        # 不加入
        tmp.pop()
        dfs(i + 1)

    dfs(0)
    return res

print(combine(4,2))
# [[0, 1], [0, 2], [0, 3], [1, 2], [1, 3], [2, 3]]