Генерация с условием через функцию 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 шт):

Автор решения: LolPopGames

Ответ был дан на основе комментария от @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
→ Ссылка
Автор решения: Stanislav Volodarskiy

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)
    )
→ Ссылка
Автор решения: CrazyElf

Вариация на ту же тему из комментария 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, вдруг кому пригодится.

→ Ссылка
Автор решения: Pak Uula

Так как 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)
→ Ссылка