Lexicographical permutations algorithm

permutations.py

from itertools import chain, permutations as perm
from operator import ge, gt, le, lt


#  Algorithm for finding permutations by moving from permutation to next lexicographical permutation
#  for more info: https://www.nayuki.io/page/next-lexicographical-permutation-algorithm


def _suffix(cont, op):
	"""
	Function returns suffix of a container fulfilling condition given by operator
	:param cont: Iterable of comparable items
	:param op: operator function (e.g. greater_than: gt)
	"""
	n = len(cont) - 1
	i, j = n, n - 1
	while j > -1 and op(cont[j], cont[i]):
		i -= 1
		j -= 1
	return cont[i:]


def _successor_idx(cont, pivot, op):
	"""
	Function provides successor of a pivot in the suffix
	:param cont: Iterable of comparable items
	:param pivot: Comparable item
	:param op: operator function
	:return: Int
	"""
	for i, x in enumerate(cont):
		if op(cont[i], pivot):
			return i
	return -1


def _permutation(cont, suffix_op, successor_op):
	"""
	Next/Prev Permutation function
	:param cont: Iterable of comparable items
	:param suffix_op: operator function
	:param successor_op: operator function
	:return: cont: Iterable of comparable items
	"""
	suffix = _suffix(cont, suffix_op)[::-1]
	prefix = cont[:-len(suffix)]
	while prefix:
		succ_idx = _successor_idx(suffix, prefix[-1], successor_op)
		prefix[-1], suffix[succ_idx] = suffix[succ_idx], prefix[-1]
		cont = prefix + suffix
		yield cont
		suffix = _suffix(cont, suffix_op)[::-1]
		prefix = cont[:-len(suffix)]


def _next_perm(cont):
	"""
	Function returns generator for 'next' lexicographical permutations
	:param cont: Iterable of comparable items
	"""
	return _permutation(cont, ge, gt)


def _prev_perm(cont):
	"""
	Function returns generator for 'previous' lexicographical permutations
	:param cont: Iterable of comparable items
	"""
	return _permutation(cont, le, lt)


def permutations(cont):
	"""
	Function yields all permutations of a given container
	:param cont: Iterable of comparable items
	"""
	yield cont
	for p in chain(_next_perm(cont), _prev_perm(cont)):
		yield p


def test():
	cs = ['abc', '1234', '4512', '0125330', '000', '1', '211', '22']
	for cont in cs:
		ms = sorted(''.join(p) for p in set(perm(cont)))
		ns = sorted(''.join(p) for p in permutations([x for x in cont]))
		if ns != ms:
			return False
	return True


if __name__ == '__main__':
	ls = [1, 2, 3]
	for p in permutations(ls):
		print(p)

#  this returns all permutations of [1, 2, 3]
#  [1, 2, 3]
#  [1, 3, 2]
#  [2, 1, 3]
#  [2, 3, 1]
#  [3, 1, 2]
#  [3, 2, 1]