Python itertools.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在每次迭代中推进最右边的元素.但现在我希望能够根据以下内容指定迭代顺序:

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)

如果shapes有三个维度,axes本来例如(0, 1, 2)或(2, 0, 1)等等,所以它不是简单地使用的问题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

嗯,这实际上是一本手动专业product.它应该更快,因为轴只重新排序一次:

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)

小测试:

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

如果没有太多的产品,上面的版本很快.对于大型结果集,以下更快,但......使用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)