How to iterate the cartesian product so top items combine first?

问题

I need to get the cartesian product of iterables, like itertools.product gives me, but for optimization reasons I want those pairs/combinations with the lowest sum of indices to appear first.

So, for example, if I have two lists, a = [1, 2, 3, 4, 5] and b = ['a', 'b', 'c', 'd', 'e'], itertools.product gives me:

>>> list(itertools.product(a, b))
[(1, 'a'), (1, 'b'), (1, 'c'), (1, 'd'), (1, 'e'), (2, 'a'), (2, 'b'), (2, 'c'), (2, 'd'), (2, 'e'), (3, 'a'), (3, 'b'), (3, 'c'), (3, 'd'), (3, 'e'), (4, 'a'), (4, 'b'), (4, 'c'), (4, 'd'), (4, 'e'), (5, 'a'), (5, 'b'), (5, 'c'), (5, 'd'), (5, 'e')]

Instead, I would want to see (2, 'a') before (1, 'c'). The exact order, between e.g. (1, 'b') and (2, 'a'), is unimportant.

Currently, I am sorting a list based on the product of the index ranges:

>>> sorted(list(itertools.product(range(len(a)), range(len(b)))), lambda a, b: sum(a) - sum(b))
[(0, 0), (0, 1), (1, 0), (0, 2), (1, 1), (2, 0), (0, 3), (1, 2), (2, 1), (3, 0), (0, 4), (1, 3), (2, 2), (3, 1), (4, 0), (1, 4), (2, 3), (3, 2), (4, 1), (2, 4), (3, 3), (4, 2), (3, 4), (4, 3), (4, 4)]

Then using that to index the lists. However, this takes too much memory with long lists. I need some kind of generator with the same calling convention as itertools.product, but I cannot figure out the way to iterate so that I get both the ordering and all the possible pairs exactly once.

回答1:

def cartprod(x,y):
    nx = len(x)
    ny = len(y)
    for i in range(nx+ny):
        for j in range(max(0,i-ny+1), min(i+1,nx)):
            yield (x[j],y[i-j])

回答2:

updated following @otus comment - generating indices ordered by sum, using those to lookup values:

A = range(5)
B = 'abcde'

def indices(A,B):
    # iterate all possible target sums in order
    for m in range(max(A)+max(B)):
        for a in A:
            # stop once current target sum isn't possible
            if a > m:
                break
            # yield if sum equals current target sum
            if m-a in B:
                yield a,m-a

def values(A,B):
    for a,b in indices(range(len(A)),set(range(len(B)))):
        yield A[a],B[b]

print list(values(A,B))

output:

[(0, 'a'), (0, 'b'), (1, 'a'), (0, 'c'), (1, 'b'), (2, 'a'), (0, 'd'), (1, 'c'), (2, 'b'), (3, 'a'), (0, 'e'), (1, 'd'), (2, 'c'), (3, 'b'), (4, 'a'), (1, 'e'), (2, 'd'), (3, 'c'), (4, 'b'), (2, 'e'), (3, 'd'), (4, 'c'), (3, 'e'), (4, 'd')]

来源：https://stackoverflow.com/questions/32603278/how-to-iterate-the-cartesian-product-so-top-items-combine-first

标签

python

iteration

itertools