mmj - 1 year ago 53

Python Question

Given a combination of

`k`

`n`

`itertools.combination(range(1,n),k)`

`list`

`dict`

Since

`itertools`

By the way here is my solution:

`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

where

`nck`

For example:

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

find_idx(comb,14) # -> 654

And here is an equivalent but maybe less involved version (actually I derived the previous one from the following one). I considered the integers of the combination

`c`

`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

Answer Source

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.

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):
...
```