tlorin
tlorin

Reputation: 1150

Python getting unique pairs from multiple lists of different lengths

Suppose that in Python I have 3 lists: a, b, c of variable lengths. For example :

a=[1,2,3]
b=[4,5,6]
c=[7,8]

I would like to get every unique combination of TWO elements of the 3 lists above, i. e.

[1,4],[1,5],[1,6],[1,7],[1,8],[2,4],[2,5]... and NOT unique combinations of the 3 lists (such as [1,4,7],[1,4,8],...).

I have looked at the solution here using itertools that is perfectly fine for 2 lists ; however, this solution does not work anymore when including a nth list because the unique combinations are of length n.

Here is what I have tried:

import itertools

a=[1,2,3]
b=[4,5,6]
c=[7,8]

d=list(itertools.product(a,b,c))

[(1, 4, 7), (1, 4, 8), (1, 5, 7), (1, 5, 8), (1, 6, 7), (1, 6, 8), (2, 4, 7), (2, 4, 8), (2, 5, 7), (2, 5, 8), (2, 6, 7), (2, 6, 8), (3, 4, 7), (3, 4, 8), (3, 5, 7), (3, 5, 8), (3, 6, 7), (3, 6, 8)]

Note: Above is just an example and the solution should work for n lists of variable length and with possibly the same value being in different lists... Any idea of how I could do would be greatly appreciated! :)


EDIT: as asked by @SirParselot, the elements have to come from different lists

Upvotes: 7

Views: 6668

Answers (2)

PM 2Ring
PM 2Ring

Reputation: 55499

You want the Cartesian product of each pair of lists in (a, b, c), so first you need itertools.combinations to generate the pairs of lists, and then itertools.product to create the desired output tuples.

from itertools import combinations, product

def pairs(*lists):
    for t in combinations(lists, 2):
        for pair in product(*t):
            yield pair

a = [1, 2, 3]
b = [4, 5, 6]
c = [7, 8]

for pair in pairs(a, b, c):
    print(pair)

output

(1, 4)
(1, 5)
(1, 6)
(2, 4)
(2, 5)
(2, 6)
(3, 4)
(3, 5)
(3, 6)
(1, 7)
(1, 8)
(2, 7)
(2, 8)
(3, 7)
(3, 8)
(4, 7)
(4, 8)
(5, 7)
(5, 8)
(6, 7)
(6, 8)

Here's a new version that handles repeated elements. It does not return a tuple if the two items in the tuple equal each other, and it also eliminates duplicated tuples in the output by feeding the output from pairs into a set. This is reasonably efficient since pairs is a generator, so duplicates are eliminated as they are found.

from itertools import combinations, product

def pairs(*lists):
    for t in combinations(lists, 2):
        for pair in product(*t):
            #Don't output pairs containing duplicated elements 
            if pair[0] != pair[1]:
                yield pair

a = [1, 2, 3]
b = [3, 4, 5]
c = [5, 6]

#Feed the output of `pairs` into a set to eliminate duplicate tuples
output = set(pairs(a, b, c))
for pair in sorted(output):
    print(pair)

output

(1, 3)
(1, 4)
(1, 5)
(1, 6)
(2, 3)
(2, 4)
(2, 5)
(2, 6)
(3, 4)
(3, 5)
(3, 6)
(4, 5)
(4, 6)
(5, 6)

Upvotes: 10

DainDwarf
DainDwarf

Reputation: 1669

I guess what you should do is use your working solution for two lists to do it with n lists. Basically, you could transform your input into a list of lists, and do:

for index, left_list in enumerate(list_of_lists):
        other_lists = (list_of_lists[0:index] + list_of_lists[index+1:])
        right_list = [i for sublist in other_lists for i in sublist]
        print(list(itertools.product(left_list, right_list)))

Upvotes: 0

Related Questions