Генерация с условием через функцию itertools.permutations()
Возник вопрос о генерации с условиями с использованием функции permutations
из библиотеки itertools:
Имеется следующий код:
combinations = list(itertools.permutations([1, 2, 3], 3))
В результате его выполнения получаем:
(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)
Это все возможные вариации сочетания заданных чисел. Возможно ли генерировать (именно генерировать, а не показывать нужные вариации среди уже всех сгенерированных) конкретные вариации с указанием условия, на какую цифру они должны начинаться? То есть что-то наподобие вот этого:
combinations = list(itertools.permutations([1, 2, 3], 3, only 1)
Результат:
(1, 2, 3), (1, 3, 2)
Ответы (4 шт):
Ответ был дан на основе комментария от @aleksandrbarakin. Он обозначен флажком "Общий".
# Суть алга - убирать первый элемент,
# Генерировать без него,
# А затем результаты добовлять этот элемент ко всем сгенерированным результатам
def permutations(variants: list, v_len: int, first_element) -> list:
variants.remove(first_element) # Удаляем "first_element"
combinations = list(itertools.permutations(variants, v_len-1)) # Вычисляем
combinations = [[first_element] + list(x) for x in combinations]
return combinations
NB Идея уже рассказана в комментариях. Я только хочу привести правильную, по моему мнению, реализацию.
Исключаем "первое значение" из списка, если оно там было. Генерируем перестановки без него, добавляем его в начало каждой перестановки:
def limited_permutations(iterable, first_value, r=None):
values = list(iterable)
if first_value in values:
i = values.index(first_value)
values = values[:i] + values[i + 1:]
if r is not None:
r -= 1
return (
(first_value, ) + p
for p in itertools.permutations(values, r=r)
)
Вариация на ту же тему из комментария
aleksandr barakin, чисто на функциях из itertools
для красоты, без проверок, считаем, что функция уже вызвана нужным образом с отделением числа от списка заранее, ну и генератор в генераторе получается, для экономии памяти и экономного досрочного завершения при необходимости:
from itertools import permutations, cycle, chain
def combinations(first, lst):
for prefix, body in zip(cycle([first]), permutations(lst, len(lst))):
yield chain([prefix], body)
for item in combinations(1, [2, 3]):
print(list(item))
Вывод:
[1, 2, 3]
[1, 3, 2]
Update: Лишнего я навертел, всё ещё проще:
from itertools import permutations, chain
def combinations(first, lst):
return (chain([first], body) for body in permutations(lst, len(lst)))
А это уже практически сводится к ответу Stanislav Volodarskiy, только без проверок. Ну ладно, оставлю только как пример применения cycle
и chain
, вдруг кому пригодится.
Так как itertools не позволяет влезть внутрь генератора и добавить проверки, то нужно делать свой.
Вот как можно сделать генератор перестановок с проверкой частичной перестановки.
Генератор строит перестановки рекурсивно:
- Сначала префикс пуст
- Если длина префикса равна длине набора элементов, вернуть префикс - перестановка построена
- Если префикс меньше набора элементов, то вернуть все перестановки, начинающиеся с данного префикса, добавляя к префиксу поочередно оставшиеся элементы и вызывая генератор с удлинённым префиксом.
Предикат получает префикс перестановки и возвращает одно из трёх значений:
- OH_SHIT: префикс нарушает ограничение, не нужно строить остальные перестановки, которые начинаются с данного префикса.
- SHOOT_IT_BABY: все перестановки, начинающиеся с данного префикса, заведомо корректны. Дочерние префиксы проверять не нужно.
- SO_FAR_SO_GOOD: префикс корректен, но дочерние префиксы тоже нужно проверять.
Начну с примеров:
Все перестановки четырёх элементов:
p_gen = permute(list("ABCD")) print(list(p_gen))
[('A', 'B', 'C', 'D'), ('A', 'B', 'D', 'C'), ('A', 'C', 'B', 'D'), ('A', 'C', 'D', 'B'), ('A', 'D', 'B', 'C'), ('A', 'D', 'C', 'B'), ('B', 'A', 'C', 'D'), ('B', 'A', 'D', 'C'), ('B', 'C', 'A', 'D'), ('B', 'C', 'D', 'A'), ('B', 'D', 'A', 'C'), ('B', 'D', 'C', 'A'), ('C', 'A', 'B', 'D'), ('C', 'A', 'D', 'B'), ('C', 'B', 'A', 'D'), ('C', 'B', 'D', 'A'), ('C', 'D', 'A', 'B'), ('C', 'D', 'B', 'A'), ('D', 'A', 'B', 'C'), ('D', 'A', 'C', 'B'), ('D', 'B', 'A', 'C'), ('D', 'B', 'C', 'A'), ('D', 'C', 'A', 'B'), ('D', 'C', 'B', 'A')]
Все перестановки, которые не начинаются с 'A'
p_gen = permute(list("ABCD"), lambda lst: OH_SHIT if lst[0] == "A" else SHOOT_IT_BABY) print(list(p_gen))
[('B', 'A', 'C', 'D'), ('B', 'A', 'D', 'C'), ('B', 'C', 'A', 'D'), ('B', 'C', 'D', 'A'), ('B', 'D', 'A', 'C'), ('B', 'D', 'C', 'A'), ('C', 'A', 'B', 'D'), ('C', 'A', 'D', 'B'), ('C', 'B', 'A', 'D'), ('C', 'B', 'D', 'A'), ('C', 'D', 'A', 'B'), ('C', 'D', 'B', 'A'), ('D', 'A', 'B', 'C'), ('D', 'A', 'C', 'B'), ('D', 'B', 'A', 'C'), ('D', 'B', 'C', 'A'), ('D', 'C', 'A', 'B'), ('D', 'C', 'B', 'A')]
Все перестановки, в которых 'B' предшествует 'C'.
p_gen = permute( list("ABCD"), lambda lst: OH_SHIT if 'C' in lst else SHOOT_IT_BABY if 'B' in lst else SO_FAR_SO_GOOD ) print(list(p_gen))
[('A', 'B', 'C', 'D'), ('A', 'B', 'D', 'C'), ('A', 'D', 'B', 'C'), ('B', 'A', 'C', 'D'), ('B', 'A', 'D', 'C'), ('B', 'C', 'A', 'D'), ('B', 'C', 'D', 'A'), ('B', 'D', 'A', 'C'), ('B', 'D', 'C', 'A'), ('D', 'A', 'B', 'C'), ('D', 'B', 'A', 'C'), ('D', 'B', 'C', 'A')]
Кроме перестановок, которые заканчиваются на 'C'
p_gen = permute( list("ABCD"), lambda lst: OH_SHIT if (len(lst) == 4 and 'C' == lst[-1]) else SO_FAR_SO_GOOD ) print(list(p_gen))
[('A', 'B', 'C', 'D'), ('A', 'C', 'B', 'D'), ('A', 'C', 'D', 'B'), ('A', 'D', 'C', 'B'), ('B', 'A', 'C', 'D'), ('B', 'C', 'A', 'D'), ('B', 'C', 'D', 'A'), ('B', 'D', 'C', 'A'), ('C', 'A', 'B', 'D'), ('C', 'A', 'D', 'B'), ('C', 'B', 'A', 'D'), ('C', 'B', 'D', 'A'), ('C', 'D', 'A', 'B'), ('C', 'D', 'B', 'A'), ('D', 'A', 'C', 'B'), ('D', 'B', 'C', 'A'), ('D', 'C', 'A', 'B'), ('D', 'C', 'B', 'A')]
Перестановщик
OH_SHIT = 0
SO_FAR_SO_GOOD = 1
SHOOT_IT_BABY = 2
def _assert_verdict(v: int) -> None:
if v == OH_SHIT or v == SO_FAR_SO_GOOD or v == SHOOT_IT_BABY:
return
raise ValueError(f"Invalid verdict: {v}")
def _all_is_good(_):
return SO_FAR_SO_GOOD
class _Permutator[T]:
"""Строит перестановки заданных элементов в соответствии с предикатом.
Предикат получает префикс перестановки и возвращает одно из трёх значений:
- OH_SHIT: префикс нарушает ограничение, не нужно строить остальные перестановки,
которые начинаются с данного префикса.
- SHOOT_IT_BABY: все перестановки, начинающиеся с данного префикса, заведомо корректны.
Дочерние префиксы проверять не нужно.
- SO_FAR_SO_GOOD: префикс корректен, но дочерние префиксы тоже нужно проверять.
Объект типа `_Permutator` является итерируемым. Обратный итератор не определён.
"""
def __init__(self, lst: list[T], predicate: callable = _all_is_good):
self.lst = lst
self.predicate = predicate
# Состояние перестановщика:
# - текущий префикс перестановки
self.prefix: List[T] = []
# - индексы элементов, включенных в префикс
self.idx_set : set[int] = set()
# - состояние предиката
self.check_state = SO_FAR_SO_GOOD
def _mk_permut_recursive(self) -> Iterator[Tuple[T]]:
"""Рекурсивно строит все перестановки, удовлетворяющие предикату."""
if len(self.prefix) == len(self.lst):
# все элементы уже в префиксе, перестановка построена. Префикс проверен выше по стеку,
# можно возвращать кортеж элементов без проверки предиката
yield tuple(self.prefix)
return
for i in range(len(self.lst)):
# Оптимизация проверки, что элемент `lst[i]` уже включен в префикс
if i in self.idx_set:
continue
old_state = self.check_state
try:
# строим дочерний префикс
self.prefix.append(self.lst[i])
if self.check_state != SHOOT_IT_BABY:
# нужно проверить префикс
v = self.predicate(self.prefix)
_assert_verdict(v)
if v == OH_SHIT:
# префикс нарушает предикат,
# дальше перестановки с этим префиксом не строим.
continue
self.check_state = v
# префикс корректен (либо состояние SHOOT_IT_BABY, либо вердикт != OH_SHIT)
# Заносим индекс добавленного элемента в множество индексов перестановки
self.idx_set.add(i)
# возвращаем все перестановки с дочерним префиксом
yield from self._mk_permut_recursive()
finally:
# Восстанавливаем состояние перестановщика
self.prefix.pop()
self.check_state = old_state
self.idx_set.discard(i)
def __iter__(self):
return self._mk_permut_recursive()
def __call__(self):
return iter(self)
и функция-обёртка
def permute[T](elems: list[T]|tuple[T], predicate: callable = _all_is_good) -> Iterator[Tuple[T]]:
"""Возвращает генератор перестановок, удовлетворяющих предикату.
Предикат получает префикс перестановки и возвращает одно из трёх значений:
- OH_SHIT: префикс нарушает ограничение, не нужно строить остальные перестановки,
которые начинаются с данного префикса.
- SHOOT_IT_BABY: все перестановки, начинающиеся с данного префикса, заведомо корректны.
Дочерние префиксы проверять не нужно.
- SO_FAR_SO_GOOD: префикс корректен, но дочерние префиксы тоже нужно проверять.
"""
return _Permutator(elems, predicate=predicate)