Python itertools.combinations 和 itertools.permutations 等价代码实现

时间:2022-05-07 08:44:39

最近编程时经常要用到排序组合的代码,想当年还抱着一些情况买了一本《组合数学》,不过现在这货也不知道被自己放哪里了,估计不会是垫桌子腿了吧。

由于去年去东北大学考博面试的时候遇到过可能涉及排列组合的编程题,要我当时很是十分的下不来台,所以现在遇到这货比较敏感,虽然用到的时候直接调用库就万事大吉,不过细细想来总觉不妥,遂把库函数的说明文档调了出来。

https://docs.python.org/dev/library/itertools.html#itertools.combinations

itertools.combinations(iterable, r)

Return r length subsequences of elements from the input iterable.

Combinations are emitted in lexicographic sort order. So, if the input iterable is sorted, the combination tuples will be produced in sorted order.

Elements are treated as unique based on their position, not on their value. So if the input elements are unique, there will be no repeat values in each combination.

Roughly equivalent to:

def combinations(iterable, r):
# combinations('ABCD', 2) --> AB AC AD BC BD CD
# combinations(range(4), 3) --> 012 013 023 123
pool = tuple(iterable)
n = len(pool)
if r > n:
return
indices = list(range(r))
yield tuple(pool[i] for i in indices)
while True:
for i in reversed(range(r)):
if indices[i] != i + n - r:
break
else:
return
indices[i] += 1
for j in range(i+1, r):
indices[j] = indices[j-1] + 1
yield tuple(pool[i] for i in indices)

初看这个代码感觉比较震惊,感觉这是什么鬼,代码量不是很大,不过实现的功能倒是要人感觉很强大。

以上代码虽然是等价代码,但是毕竟为源语言实现的,由此看看还是有些帮助的,估计真实库文件可能是使用C语言写的吧,这一点就不深入研究了。

pool = tuple(iterable) 该行代码则是将输入的可迭代对象转换为元组类型。

    n = len(pool)
if r > n:
return

该段代码则是 得到元组长度n,  如果排序数 r 大于总长n, 则无意义,因此直接返回。

indices = list(range(r))
indices 为排序的索引, 也是每次迭代返回组合的索引号。
yield tuple(pool[i] for i in indices)

返回 默认的组合。如: pool 为(0,1,2,3,4), n=5, r=3, 默认的返回就是 (0,1,2) 。
        for i in reversed(range(r)):
if indices[i] != i + n - r:
break
判断返回的组合序号 是否为 终止的序号, 即(2,3,4)。

        indices[i] += 1
返回的序号中如果那个位置的序号不等于终止序列中的对应该位置的序号则将该位置序号加一。
(0,1,2) =》 (0,1,3)
(0,1,3) =》 (0,1,4) (0,1,4) =》 (0,2, 4)

 for j in range(i+1, r):
indices[j] = indices[j-1] + 1
不同的序号位置加一后,其后续的位置依次在其前一位置上加一。即,(0,2,4)=》(0,2,3), (0,3,4)=》(1,3,4)=》(1,2,4)=》(1,2,3)

https://docs.python.org/dev/library/itertools.html#itertools.permutations
itertools.permutations(iterable, r=None)

Return successive r length permutations of elements in the iterable.

If r is not specified or is None, then r defaults to the length of the iterable and all possible full-length permutations are generated.

Permutations are emitted in lexicographic sort order. So, if the input iterable is sorted, the permutation tuples will be produced in sorted order.

Elements are treated as unique based on their position, not on their value. So if the input elements are unique, there will be no repeat values in each permutation.

Roughly equivalent to:

 
def permutations(iterable, r=None):
# permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
# permutations(range(3)) --> 012 021 102 120 201 210
pool = tuple(iterable)
n = len(pool)
r = n if r is None else r
if r > n:
return
indices = list(range(n))
cycles = list(range(n, n-r, -1))
yield tuple(pool[i] for i in indices[:r])
while n:
for i in reversed(range(r)):
cycles[i] -= 1
if cycles[i] == 0:
indices[i:] = indices[i+1:] + indices[i:i+1]
cycles[i] = n - i
else:
j = cycles[i]
indices[i], indices[-j] = indices[-j], indices[i]
yield tuple(pool[i] for i in indices[:r])
break
else:
return