Python itertools.product reorder the generation

问题

I have this:

shape = (2, 4) # arbitrary, could be 3 dimensions such as (3, 5, 7), etc...

for i in itertools.product(*(range(x) for x in shape)):
    print(i)

# output: (0, 0) (0, 1) (0, 2) (0, 3) (1, 0) (1, 1) (1, 2) (1, 3)

So far, so good, itertools.product advances the rightmost element on every iteration. But now I want to be able to specify the iteration order according to the following:

axes = (0, 1) # normal order
# output: (0, 0) (0, 1) (0, 2) (0, 3) (1, 0) (1, 1) (1, 2) (1, 3)

axes = (1, 0) # reversed order
# output: (0, 0) (1, 0) (2, 0) (3, 0) (0, 1) (1, 1) (2, 1) (3, 1)

If shapes had three dimensions, axes could have been for instance (0, 1, 2) or (2, 0, 1) etc, so it's not a matter of simply using reversed(). So I wrote some code that does that but seems very inefficient:

axes = (1, 0)

# transposed axes
tpaxes = [0]*len(axes)
for i in range(len(axes)):
    tpaxes[axes[i]] = i

for i in itertools.product(*(range(x) for x in shape)):
    # reorder the output of itertools.product
    x = (i[y] for y in tpaxes)
    print(tuple(x))

Any ideas on how to properly do this?

回答1:

Well, this is in fact a manual specialised product. It should be faster since axes are reordered only once:

def gen_chain(dest, size, idx, parent):
    # iterate over the axis once
    # then trigger the previous dimension to update
    # until everything is exhausted
    while True:
        if parent: next(parent) # StopIterator is propagated upwards

        for i in xrange(size):
            dest[idx] = i
            yield 

        if not parent: break

def prod(shape, axes):
    buf = [0] * len(shape)
    gen = None

    # EDIT: fixed the axes order to be compliant with the example in OP 
    for s, a in zip(shape, axes):
        # iterate over the axis and put to transposed
        gen = gen_chain(buf, s, a, gen)

    for _ in gen:
        yield tuple(buf)


print list(prod((2,4), (0,1)))
# [(0, 0), (0, 1), (0, 2), (0, 3), (1, 0), (1, 1), (1, 2), (1, 3)]
print list(prod((2,4), (1,0)))
# [(0, 0), (1, 0), (2, 0), (3, 0), (0, 1), (1, 1), (2, 1), (3, 1)]
print list(prod((4,3,2),(1,2,0)))
# [(0, 0, 0), (1, 0, 0), (0, 0, 1), (1, 0, 1), (0, 0, 2), (1, 0, 2), ...

回答2:

If you can afford it memory-wise: Let itertools.product do the hard work, and use zip to switch the axes around.

import itertools
def product(shape, axes):
    prod_trans = tuple(zip(*itertools.product(*(range(shape[axis]) for axis in axes))))

    prod_trans_ordered = [None] * len(axes)
    for i, axis in enumerate(axes):
        prod_trans_ordered[axis] = prod_trans[i]
    return zip(*prod_trans_ordered)

Small test:

>>> print(*product((2, 2, 4), (1, 2, 0)))
(0, 0, 0) (1, 0, 0) (0, 0, 1) (1, 0, 1) (0, 0, 2) (1, 0, 2) (0, 0, 3) (1, 0, 3) (0, 1, 0) (1, 1, 0) (0, 1, 1) (1, 1, 1) (0, 1, 2) (1, 1, 2) (0, 1, 3) (1, 1, 3)

The above version is fast if there are not too may products. For large result sets, the following is faster, but... uses eval (although in a rather safe way):

def product(shape, axes):
    d = dict(("r%i" % axis, range(shape[axis])) for axis in axes)
    text_tuple = "".join("x%i, " % i for i in range(len(axes)))
    text_for = " ".join("for x%i in r%i" % (axis, axis) for axis in axes)
    return eval("((%s) %s)" % (text_tuple, text_for), d)

Edit: If you want to not only change the order of iteration, but also the shape (as in the OP's example), small changes are needed:

import itertools
def product(shape, axes):
    prod_trans = tuple(zip(*itertools.product(*(range(s) for s in shape))))

    prod_trans_ordered = [None] * len(axes)
    for i, axis in enumerate(axes):
        prod_trans_ordered[axis] = prod_trans[i]
    return zip(*prod_trans_ordered)

And the eval version:

def product(shape, axes):
    d = dict(("r%i" % axis, range(s)) for axis, s in zip(axes, shape))
    text_tuple = "".join("x%i, " % i for i in range(len(axes)))
    text_for = " ".join("for x%i in r%i" % (axis, axis) for axis in axes)
    return eval("((%s) %s)" % (text_tuple, text_for), d)

Test:

>>> print(*product((2, 2, 4), (1, 2, 0)))
(0, 0, 0) (1, 0, 0) (2, 0, 0) (3, 0, 0) (0, 0, 1) (1, 0, 1) (2, 0, 1) (3, 0, 1) (0, 1, 0) (1, 1, 0) (2, 1, 0) (3, 1, 0) (0, 1, 1) (1, 1, 1) (2, 1, 1) (3, 1, 1)

回答3:

I don't know how efficient this is, but you should be able to do something like this...

shape = (2, 4, 3)
axes = (2, 0, 1)

# Needed to get the original ordering back
axes_undo = tuple(reversed(axes))

# Reorder the shape in a configuration so that .product will give you
# the order you want.
reordered = tuple(reversed(map(lambda x: shape[x], list(axes))))

# When printing out the results from .product, put the results back
# into the original order.
for i in itertools.product(*(range(x) for x in reordered)):
    print(tuple(map(lambda x: i[x], list(axes_undo))))

I tried is up to 4 dimensions and it seems to work. ;)

I'm just swapping the dimensions around and then swapping them back.

回答4:

Have you tried timing to see how much longer it takes? What you have shouldn't be much slower than without reordering.

You could try modify what you have to use in-place splice assignment.

tpaxes = tuple(tpaxes)
for i in itertools.product(*(range(x) for x in shape)):
    # reorder the output of itertools.product
    i[:] = (i[y] for y in tpaxes)
    print(tuple(x))

Also you could get a speedup by making tpaxes a local variable of a function rather than a global variable (which has slower lookup times)

Otherwise my suggestion is somehow write your own product function..

回答5:

for i in itertools.product(*(range(x) for x in reversed(shape))):
    print tuple(reversed(i))

回答6:

import itertools

normal = (0, 1)
reverse = (1, 0)

def axes_ordering(x):
    a, b = x
    return b - a

shape = (2, 4)

for each in itertools.product(*(range(x) for x in shape)):
    print(each[::axes_ordering(normal)], each[::axes_ordering(reverse)])

result:

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

来源：https://stackoverflow.com/questions/9903257/python-itertools-product-reorder-the-generation

标签

python

itertools