Pythonitertools.product重新排序生成

Question

我有这个：

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)

到目前为止，很好，itertools.product在每次迭代中推进最右边的元素。但现在我希望能够根据以下内容指定迭代顺序：

^{pr2}$

如果shapes有三个维度，axes可能是(0, 1, 2)或{}等，所以这不是简单使用reversed()的问题。所以我写了一些代码，但是看起来效率很低：

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))

有什么好办法吗？在

Answer 1

实际上，这是一本专门的手册。由于轴只重新排序一次，所以速度应该更快：

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), ...

Answer 2

如果你有足够的内存空间：让itertools.product完成这项艰苦的工作，并使用zip来切换轴。在

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)

小测试：

^{2}$

如果没有太多的产品，以上版本是快速的。对于大的结果集，下面的方法更快，但是。。。使用eval（尽管以一种相当安全的方式）：

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)

编辑：如果您不仅要更改迭代的顺序，还需要更改形状（如OP的示例中所示），则需要进行小的更改：

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)

以及eval版本：

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)

测试：

>>> 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)

Answer 3

我不知道这有多有效，但你应该能做这样的事。。。在

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))))

我尝试了4个维度，它似乎有效。；）

我只是交换一下尺寸然后再交换回来。在