Find the index of a given combination (of natural numbers) among those returned by `itertools` Python module

前端 未结 3 1035
无人共我
无人共我 2021-01-05 08:01

Given a combination of k of the first n natural numbers, for some reason I need to find the position of such combination among those returned by

3条回答
  •  一生所求
    2021-01-05 08:44

    Your solution seems quite fast. In find_idx, you have two for loop, the inner loop can be optimized using the formular:

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

    so, you can replace sum(nck(n-2-x,k-1) for x in range(c-last_c-1)) with nck(n-1, k) - nck(n-c+last_c, k).

    I don't know how you implement your nck(n, k) function, but it should be O(k) measured in time complexity. Here I provide my implementation:

    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)
    

    Finally, your solution become O(k^2) without recursion. It's quite fast since k wouldn't be too large.

    Update

    I've noticed that nck's parameters are (n, k). Both n and k won't be too large. We may speed up the program by caching.

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

    In python3 this can be done by using functools.lru_cache decorator:

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

提交回复
热议问题