Get r-length tuple combinations of m lists, with no more than a single element from any list, and r < m

Question

In the example below, I have m = 3 lists which and I calculate the combinations of size r = 2 .

import itertools

a = ['a1', 'a2', 'a3']
b = ['b1', 'b2', 'b3']
c = ['c1', 'c2', 'c3']

print(list(itertools.combinations(itertools.chain(a, b, c), 2)))

Output:

[('a1', 'a2'), ('a1', 'a3'), ('a1', 'b1'), ('a1', 'b2'), ('a1', 'b3'), ('a1', 'c1'), ('a1', 'c2'), ('a1', 'c3'), ('a2', 'a3'), ('a2', 'b1'), ('a2', 'b2'), ('a2', 'b3'), ('a2', 'c1'), ('a2', 'c2'), ('a2', 'c3'), ('a3', 'b1'), ('a3', 'b2'), ('a3', 'b3'), ('a3', 'c1'), ('a3', 'c2'), ('a3', 'c3'), ('b1', 'b2'), ('b1', 'b3'), ('b1', 'c1'), ('b1', 'c2'), ('b1', 'c3'), ('b2', 'b3'), ('b2', 'c1'), ('b2', 'c2'), ('b2', 'c3'), ('b3', 'c1'), ('b3', 'c2'), ('b3', 'c3'), ('c1', 'c2'), ('c1', 'c3'), ('c2', 'c3')]

Problem:

I don't want the combinations that come from the same list. For example, ('a1', 'a2') and ('a1', 'a3') are should be removed.

Answer 1

To slightly paraphrase the answer by Mateen, you could first choose r sets in m choose r ways in the outer loop, then iterate over the product of r sets in the inner loop. The order of the output may be different from other approaches.

lsts = [a, b, c]
def f(lsts, r):
  """
  lsts :: [[a]]
  r :: Integer
  Generate r-tuples with at most one element coming from
  each member of lsts.
  """
  m = len(lsts)
  assert m >= r
  for xs in itertools.combinations(lsts, r):
    for x in itertools.product(*xs):
      yield x

Answer 2

This may be a go-around solution but you may want to modify the list you get as output to keep only the desired values.

[i for i in a if i[0][0]!=i[1][0]]

With a being your list a = list(itertools.combinations(itertools.chain(a,b,c), 2))

Answer 3

Filtering the output of itertools.combinations

This is not necessarily the most elegant or performant solution, but it works.

def is_unique(comb, lookup):
    xs = [lookup[x] for x in comb]
    return len(xs) == len(set(xs))


def combine(args, r):
    lookup = {x: i for i, xs in enumerate(args) for x in xs}
    return (
        comb
        for comb in itertools.combinations(itertools.chain(*args), r) 
        if is_unique(comb, lookup)
    )


>>> list(combine((a, b, c), 2))
[('a1', 'b1'), ('a1', 'b2'), ('a1', 'b3'),      # (a1, b*)
 ('a1', 'c1'), ('a1', 'c2'), ('a1', 'c3'),      # (a1, c*)
 ('a2', 'b1'), ('a2', 'b2'), ('a2', 'b3'),      # (a2, b*)
 ('a2', 'c1'), ('a2', 'c2'), ('a2', 'c3'),      # (a2, c*)
 ('a3', 'b1'), ('a3', 'b2'), ('a3', 'b3'),      # (a3, b*)
 ('a3', 'c1'), ('a3', 'c2'), ('a3', 'c3'),      # (a3, c*)
 ('b1', 'c1'), ('b1', 'c2'), ('b1', 'c3'),      # (b1, c*)
 ('b2', 'c1'), ('b2', 'c2'), ('b2', 'c3'),      # (b2, c*)
 ('b3', 'c1'), ('b3', 'c2'), ('b3', 'c3')]      # (b3, c*)