python - 在 `itertools` Python模块返回的索引中找到给定组合(自然数)的索引

给定第一个k自然数的n组合，出于某种原因，我需要在itertools.combination(range(1,n),k)返回的组合中找到这种组合的位置(原因是这样，我可以使用list而不是dict来访问与每个组合，知道组合)。

由于itertools以常规模式产生其组合，因此可以做到这一点(而且我还找到了一种整洁的算法)，但是我正在寻找一种更快或更自然的方式，但我可能会忽略它。

顺便说一下，这是我的解决方案:

def find_idx(comb,n):
    k=len(comb)
    idx=0
    last_c=0
    for c in comb:
        #idx+=sum(nck(n-2-x,k-1) for x in range(c-last_c-1)) # a little faster without nck caching
        idx+=nck(n-1,k)-nck(n-c+last_c,k) # more elegant (thanks to Ray), faster with nck caching
        n-=c-last_c
        k-=1
        last_c=c
    return idx

其中nck返回binomial coefficient of n,k。

例如:

comb=list(itertools.combinations(range(1,14),6))[654] #pick the 654th combination
find_idx(comb,14) # -> 654

这是一个等效的但可能涉及较少的版本(实际上，我从下一个版本派生了上一个版本)。我将组合c的整数视为二进制数中1s的位置，我在解析0/1时构建了一个二叉树，并且在解析期间发现索引递增的规则模式:

def find_idx(comb,n):
    k=len(comb)
    b=bin(sum(1<<(x-1) for x in comb))[2:]
    idx=0
    for s in b[::-1]:
        if s=='0':
            idx+=nck(n-2,k-1)
        else:
            k-=1
        n-=1
    return idx

最佳答案

您的解决方案似乎很快。在find_idx中，您有两个for循环，可以使用公式来优化内部循环:

C(n, k) + C(n-1, k) + ... + C(n-r, k) = C(n+1, k+1) - C(n-r, k+1)

因此，您可以将sum(nck(n-2-x,k-1) for x in range(c-last_c-1))替换为nck(n-1, k) - nck(n-c+last_c, k)。

我不知道您如何实现nck(n, k)函数，但是应该以时间复杂度来度量O(k)。在这里，我提供了我的实现:

from operator import mul
from functools import reduce # In python 3
def nck_safe(n, k):
    if k < 0 or n < k: return 0
    return reduce(mul, range(n, n-k, -1), 1) // reduce(mul, range(1, k+1), 1)

最后，您的解决方案变为O(k ^ 2)，而无需递归。由于k不会太大，因此速度非常快。

更新

我注意到nck的参数是(n, k)。 n和k都不会太大。我们可以通过缓存来加速程序。

def nck(n, k, _cache={}):
    if (n, k) in _cache: return _cache[n, k]
    ....
    # before returning the result
    _cache[n, k] = result
    return result

在python3中，可以使用functools.lru_cache装饰器完成此操作:

@functools.lru_cache(maxsize=500)
def nck(n, k):
    ...

关于python - 在 `itertools` Python模块返回的索引中找到给定组合(自然数)的索引，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/14455634/