Search code examples
pythonpython-itertools

Python get all sets of N pairs from two lists


This is some kind of elementary task in itertools probably, but my brain is not working today and I can't seem to find it answered here already.

Suppose I have two lists with different numbers of elements. Their indices are

i1 = list(range(N))
i2 = list(range(M))

for say N=3, M=4 we have

[0, 1, 2]
[0, 1, 2, 3]

I now want to construct all possible sets of pairs of length K from these two lists. So for example for K=1 we have

(0, 0)
(0, 1)
(0, 2)
(0, 3)
(1, 0)
(1, 1)
(1, 2)
(1, 3)
(2, 0)
(2, 1)
(2, 2)
(2, 3)

and for K=2 we have

(0,0), (1,1)
(0,0), (1,2)
(0,0), (1,3)
(0,0), (2,1)
(0,0), (2,2)
(0,0), (2,3)
(0,1), (1,0)
(0,1), (1,2)
(0,1), (1,3)
(0,2), (1,0)
(0,2), (1,1)
(0,2), (1,3)
...
<and so it goes on...>

i.e. we are drawing one value from each list to make a pair, without replacement, and doing this K times to make K pairs. And we want to obtain all possible sets of such pairs.

For K=1 I guess this is the cartesian product of the lists, but it's more complicated for K=2 and above.

Note that the elements in each pair in each set are drawn from the lists without replacement. So each item from each list can appear in at most one pair in each set of pairs.

Edit: The answer here: https://stackoverflow.com/a/12935562/1447953 solves the problem when K=min(N,M), i.e. when finding the longest possible sets of pairs, but it doesn't work for smaller K .

For additional context: I am working on a spectral peak matching problem. I have two line spectra, which have varying numbers of peaks in them, and I am working on an algorithm to determine whether any of the peaks line up with each other. This involves comparing peaks from one spectrum against those in the other spectrum, so I need to iterate over all the ways the peaks might possibly be paired up for comparison (for differing numbers of potential simultaneous matches).


Solution

  • This is basically a product of k-permutations of the two lists:

    from itertools import product, permutations
    
    def f(i1, i2, k):
        for p in product(permutations(i1, k), permutations(i2, k)):
            yield tuple(zip(*p))
    

    so that:

    i1 = list(range(3))
    i2 = list(range(4))
    
    print(*f(i1, i2, 1), sep='\n')
    print(*f(i1, i2, 2), sep='\n')
    

    outputs:

    ((0, 0),)
    ((0, 1),)
    ((0, 2),)
    ((0, 3),)
    ((1, 0),)
    ((1, 1),)
    ((1, 2),)
    ((1, 3),)
    ((2, 0),)
    ((2, 1),)
    ((2, 2),)
    ((2, 3),)
    ((0, 0), (1, 1))
    ((0, 0), (1, 2))
    ((0, 0), (1, 3))
    ((0, 1), (1, 0))
    ((0, 1), (1, 2))
    ((0, 1), (1, 3))
    ((0, 2), (1, 0))
    ((0, 2), (1, 1))
    ((0, 2), (1, 3))
    ((0, 3), (1, 0))
    ((0, 3), (1, 1))
    ((0, 3), (1, 2))
    ((0, 0), (2, 1))
    ((0, 0), (2, 2))
    ((0, 0), (2, 3))
    ((0, 1), (2, 0))
    ((0, 1), (2, 2))
    ((0, 1), (2, 3))
    ((0, 2), (2, 0))
    ((0, 2), (2, 1))
    ((0, 2), (2, 3))
    ((0, 3), (2, 0))
    ((0, 3), (2, 1))
    ((0, 3), (2, 2))
    ((1, 0), (0, 1))
    ((1, 0), (0, 2))
    ((1, 0), (0, 3))
    ((1, 1), (0, 0))
    ((1, 1), (0, 2))
    ((1, 1), (0, 3))
    ((1, 2), (0, 0))
    ((1, 2), (0, 1))
    ((1, 2), (0, 3))
    ((1, 3), (0, 0))
    ((1, 3), (0, 1))
    ((1, 3), (0, 2))
    ((1, 0), (2, 1))
    ((1, 0), (2, 2))
    ((1, 0), (2, 3))
    ((1, 1), (2, 0))
    ((1, 1), (2, 2))
    ((1, 1), (2, 3))
    ((1, 2), (2, 0))
    ((1, 2), (2, 1))
    ((1, 2), (2, 3))
    ((1, 3), (2, 0))
    ((1, 3), (2, 1))
    ((1, 3), (2, 2))
    ((2, 0), (0, 1))
    ((2, 0), (0, 2))
    ((2, 0), (0, 3))
    ((2, 1), (0, 0))
    ((2, 1), (0, 2))
    ((2, 1), (0, 3))
    ((2, 2), (0, 0))
    ((2, 2), (0, 1))
    ((2, 2), (0, 3))
    ((2, 3), (0, 0))
    ((2, 3), (0, 1))
    ((2, 3), (0, 2))
    ((2, 0), (1, 1))
    ((2, 0), (1, 2))
    ((2, 0), (1, 3))
    ((2, 1), (1, 0))
    ((2, 1), (1, 2))
    ((2, 1), (1, 3))
    ((2, 2), (1, 0))
    ((2, 2), (1, 1))
    ((2, 2), (1, 3))
    ((2, 3), (1, 0))
    ((2, 3), (1, 1))
    ((2, 3), (1, 2))
    

    Demo: https://ideone.com/WVkVKW