Combinações de uma lista em Python

Use itertools.combinations() para selecionar uma quantidade fixa de itens sem reutilizar uma posição da entrada. Se for permitido selecionar uma posição mais de uma vez, use itertools.combinations_with_replacement(). Ambas as funções geram combinações, nas quais a ordem da seleção não cria um novo resultado; nas permutações, ordens diferentes são consideradas distintas.

Escolha o tipo certo de combinação

Escolha um tamanho fixo para a seleção com combinations(), permita repetir uma posição com combinations_with_replacement() ou gere todos os subconjuntos com uma função powerset(). As funções da biblioteca padrão retornam iteradores de tuplas. A entrada é processada por posição, portanto valores repetidos em posições diferentes podem gerar tuplas aparentemente duplicadas. A ordem dos resultados segue a ordem das posições na entrada; os valores dentro de cada tupla não são ordenados automaticamente. Para aplicar uma regra personalizada, a recursão permite explicitar o processo de seleção.

Use itertools.combinations() para seleções sem reposição

itertools.combinations(iterable, r) produz tuplas com r elementos e não reutiliza uma posição. O iterável é percorrido na ordem em que foi fornecido. Converta o iterador em uma lista se precisar armazenar ou examinar todos os resultados de uma vez; caso contrário, percorra-o diretamente.

Execute este exemplo:

O código e os rótulos ou saídas exibidos nesta seção permanecem em inglês.

from itertools import combinations

items = ["red", "blue", "green"]
for pair in combinations(items, 2):
    print(pair)

Os resultados com dois elementos usam cada posição no máximo uma vez. Para obter itens individuais, defina r como 1; para solicitar toda a entrada em uma única tupla, defina-o como o número de posições da entrada. Se r for maior que o número de posições, o iterador não terá resultados.

Use itertools.combinations_with_replacement() quando as seleções puderem se repetir

Use itertools.combinations_with_replacement(iterable, r) quando for permitido selecionar a mesma posição da entrada mais de uma vez. Essa função também produz tuplas com r elementos, seguindo a mesma regra de ordenação por posição. Por exemplo, ao escolher pares de ["red", "blue"], o resultado inclui ("red", "red") e ("blue", "blue"); combinations() não incluiria esses pares. Isso é útil quando uma escolha repetida faz sentido, como selecionar duas bolas de sorvete que podem ter o mesmo sabor. O argumento r continua determinando o tamanho de cada tupla.

Use recursão ou um gerador para uma implementação personalizada

Uma função recursiva de retrocesso pode montar os resultados em uma lista. A cada etapa, ela acrescenta um valor, faz uma chamada recursiva a partir da posição seguinte e depois remove esse valor para que o laço possa tentar outra opção. Um gerador segue a mesma lógica de seleção, mas produz cada resultado assim que ele é encontrado. Dessa forma, não mantém toda a coleção de resultados na memória, a menos que quem o chama a converta explicitamente em uma lista; isso não significa que seja inerentemente mais rápido.

As duas funções personalizadas abaixo exigem que r seja um número inteiro não negativo e geram ValueError caso contrário. Com r = 0, cada uma retorna a única tupla vazia; quando r é positivo, mas maior que o número de posições da entrada, nenhuma produz resultados. O gerador verifica o argumento quando a iteração começa, pois o corpo de uma função geradora é executado sob demanda.

O código e os rótulos ou saídas exibidos nesta seção permanecem em inglês.

def get_combinations(items, r):
    if not isinstance(r, int) or isinstance(r, bool) or r < 0:
        raise ValueError("r must be a nonnegative integer")

    results = []

    def backtrack(start, chosen):
        if len(chosen) == r:
            results.append(tuple(chosen))
            return
        for index in range(start, len(items)):
            chosen.append(items[index])
            backtrack(index + 1, chosen)
            chosen.pop()

    backtrack(0, [])
    return results


def generate_combinations(items, r):
    if not isinstance(r, int) or isinstance(r, bool) or r < 0:
        raise ValueError("r must be a nonnegative integer")
    if r == 0:
        yield ()
        return
    for index in range(len(items)):
        for suffix in generate_combinations(items[index + 1 :], r - 1):
            yield (items[index],) + suffix


items = ["red", "blue", "green"]
print(get_combinations(items, 2))
print(list(generate_combinations(items, 2)))

Para criar uma versão personalizada sob demanda, use yield quando uma seleção completa for alcançada. Em ambas as implementações, avance a próxima chamada recursiva para index + 1 para evitar reutilizar uma posição. Se houver valores duplicados, os resultados podem parecer iguais mesmo quando vieram de posições diferentes.

Gere todos os subconjuntos com uma função powerset()

Um conjunto das partes contém subconjuntos de todos os tamanhos, incluindo a tupla vazia e a entrada completa. Esta implementação combina os resultados de cada tamanho possível. chain.from_iterable() retorna um iterador; use list() somente se precisar materializar todos os subconjuntos.

O código e os rótulos ou saídas exibidos nesta seção permanecem em inglês.

from itertools import chain, combinations


def powerset(items):
    values = list(items)
    return chain.from_iterable(
        combinations(values, size) for size in range(len(values) + 1)
    )


for subset in powerset(["red", "blue"]):
    print(subset)

Para n posições na entrada, um conjunto das partes contém 2**n subconjuntos. Essa quantidade cresce rapidamente; por isso, prefira um iterador de combinações de tamanho fixo quando a tarefa exigir seleções de um único tamanho. Use o conjunto das partes quando o resultado realmente precisar incluir subconjuntos de todos os tamanhos.