Combinação de Listas em Python

Combinação de Listas: Imagina que você tem diversas listas e você quer montar todas as combinações de um número específico de elementos de cada lista.

Implementação:

class CombinationOfListsGenerator(object): def __init__(self, lists, amount_per_list): if(amount_per_list == None or lists == None): raise Exception('Invalid argument: None.') elif len(lists) < 1: raise Exception('Invalid argument: Empty list.') elif not isinstance(amount_per_list, int): if len(lists) > len(amount_per_list): raise Exception('Invalid argument: It is missing information in amount_per_list, too few amounts.') else: for l, a in zip(lists, amount_per_list): if len(l) < a: raise Exception('Invalid argument: Amount per list is bigger than size of the list.') self.lists = lists self._current = [] self.generators = [] if isinstance(amount_per_list, int): for l in self.lists: generator = CombinationGenerator(len(l), amount_per_list) self.generators.append(generator) else: for l, a in zip(self.lists, amount_per_list): generator = CombinationGenerator(len(l), a) self.generators.append(generator) self._total = 1 for generator in self.generators: self._total *= generator.total() self.reset() def total(self): return self._total def reset(self): self._current = [] for generator in self.generators[:-1]: generator.reset() generator.next() # last one is loaded in 'next' method self.generators[-1].reset() def current(self): return self._current def has_next(self): return not (not self.generators[0].has_next() and not self.generators[-1].has_next()) def _combination_of_list(self, list, indexes): combination = [] for index in indexes: combination.append(list[index]) return combination def _update_generators(self): # Atualiza anteriores i = len(self.generators) - 1 while i >= 0 and not self.generators[i].has_next(): self.generators[i].reset() self.generators[i].next() i -= 1 if i < 0: # Tratando fim: i = 0 self.generators[i].next() def _update_current(self): self._current = [] for i in range(0, len(self.generators)): indexes = self.generators[i].current() self._current.extend(self._combination_of_list(self.lists[i], indexes)) def next(self): self._update_generators() self._update_current() return self._current
Testes de Unidade:

class CombinationOfListsTests(unittest.TestCase): def test_input_accepts_int_or_list(self): CombinationOfListsGenerator([[1], [2], [3]], 1) CombinationOfListsGenerator([[1], [2], [3]], [1,1,1]) def test_invalid_input(self): self.assertRaises(Exception, CombinationOfListsGenerator, None, 1) self.assertRaises(Exception, CombinationOfListsGenerator, [1], None) self.assertRaises(Exception, CombinationOfListsGenerator, [], 1) self.assertRaises(Exception, CombinationOfListsGenerator, [[1],[2]], 2) self.assertRaises(Exception, CombinationOfListsGenerator, [[1]], 0) # len([1]) < len([[1],[2]]) self.assertRaises(Exception, CombinationOfListsGenerator, [[1],[2]], [1]) self.assertRaises(Exception, CombinationOfListsGenerator, [[1,2],[2],[3]], [2,1]) # 2 > len([2]) self.assertRaises(Exception, CombinationOfListsGenerator, [[1],[2]], [1, 2]) self.assertRaises(Exception, CombinationOfListsGenerator, [[1],[2,3]], [1, 3]) def test_1_element_of_each_list_of_size_1(self): generator = CombinationOfListsGenerator([[1], [2]], 1) self.assertEquals(1, generator.total()) self.assertEquals([1, 2], generator.next()) def test_1_element_of_each_list_of_size_2(self): generator = CombinationOfListsGenerator([[1, 2], [3, 4]], 1) self.assertEquals(4, generator.total()) self.assertEquals([1, 3], generator.next()) self.assertEquals([1, 4], generator.next()) self.assertEquals([2, 3], generator.next()) self.assertEquals([2, 4], generator.next()) def test_2_elements_of_each_list_of_size_2(self): generator = CombinationOfListsGenerator([[1, 2], [3, 4]], 2) self.assertEquals(1, generator.total()) self.assertEquals([1, 2, 3, 4], generator.next()) def test_2_elements_of_each_list_of_size_3(self): generator = CombinationOfListsGenerator([[1, 2, 3], [4, 5, 6]], 2) self.assertEquals(9, generator.total()) self.assertEquals([1, 2, 4, 5], generator.next()) self.assertEquals([1, 2, 4, 6], generator.next()) self.assertEquals([1, 2, 5, 6], generator.next()) self.assertEquals([1, 3, 4, 5], generator.next()) self.assertEquals([1, 3, 4, 6], generator.next()) self.assertEquals([1, 3, 5, 6], generator.next()) self.assertEquals([2, 3, 4, 5], generator.next()) self.assertEquals([2, 3, 4, 6], generator.next()) self.assertEquals([2, 3, 5, 6], generator.next()) def test_1_element_of_each_list_of_size_2(self): generator = CombinationOfListsGenerator([[1, 2], [3, 4], [5, 6]], 1) self.assertEquals(8, generator.total()) self.assertEquals([1, 3, 5], generator.next()) self.assertEquals([1, 3, 6], generator.next()) self.assertEquals([1, 4, 5], generator.next()) self.assertEquals([1, 4, 6], generator.next()) self.assertEquals([2, 3, 5], generator.next()) self.assertEquals([2, 3, 6], generator.next()) self.assertEquals([2, 4, 5], generator.next()) self.assertEquals([2, 4, 6], generator.next()) def test_custom_elements_of_each_list_of_size_2(self): generator = CombinationOfListsGenerator([[1, 2], [3, 4], [5, 6]], [1, 2, 1]) self.assertEquals(4, generator.total()) self.assertEquals([1, 3, 4, 5], generator.next()) self.assertEquals([1, 3, 4, 6], generator.next()) self.assertEquals([2, 3, 4, 5], generator.next()) self.assertEquals([2, 3, 4, 6], generator.next()) def test_has_next(self): generator = CombinationOfListsGenerator([[1, 2], [3, 4]], 1) self.assertEquals(4, generator.total()) self.assertTrue(generator.has_next()) self.assertEquals([1, 3], generator.next()) self.assertTrue(generator.has_next()) self.assertEquals([1, 4], generator.next()) self.assertTrue(generator.has_next()) self.assertEquals([2, 3], generator.next()) self.assertTrue(generator.has_next()) self.assertEquals([2, 4], generator.next()) self.assertFalse(generator.has_next()) generator.next() self.assertTrue(generator.has_next()) def test_reset(self): generator = CombinationOfListsGenerator([[1, 2], [3, 4]], 1) self.assertEquals(4, generator.total()) self.assertEquals([1, 3], generator.next()) generator.reset() self.assertEquals([1, 3], generator.next()) def test_current(self): generator = CombinationOfListsGenerator([[1, 2], [3, 4]], 1) self.assertEquals(4, generator.total()) self.assertEquals([], generator.current()) self.assertEquals([1, 3], generator.next()) self.assertEquals([1, 3], generator.current()) self.assertEquals([1, 4], generator.next()) self.assertEquals([1, 4], generator.current())