Python itertools.product 重新排序生成

Posted 2023-02-25

【中文标题】Python itertools.product 重新排序生成

【英文标题】：Python itertools.product reorder the generation

【发布时间】：2012-03-28 07:52:45

【问题描述】：

我有这个：

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

关于如何正确执行此操作的任何想法？

【问题讨论】：

对于您的示例，tpaxes 是 [1, 0]，axes 是 (1, 0)。为了清楚起见，您可能希望更改示例数据，以便它们有所不同:)

是的，axes=tpaxes 因为这是对二维矩阵的轴重新排序的唯一可能方法。对于 3d 矩阵，情况并非如此。例如，如果轴是 (2, 0, 1)，那么 tpaxes 将是 (1, 2, 0)。

我知道——只是想指出，在这种情况下，一个更复杂的例子会更好；没有冒犯。

唯一没有额外步骤的方法是编写自己的product 实现。我链接到一对你可以从前几天开始的this post about itertools.product。我的问题是为什么。如果您确实需要在某些特定情况下执行此操作，您只需将提供给产品的参数重新排序为正确的开始顺序，而无需更改生成的顺序。

你不能事后对product的输出进行排序吗？

【参考方案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), ...

【参考方案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)

【讨论】：

+1，我特别喜欢eval 解决方案。这是一个每天都看不到的创意编码！它可能很难看，但它的作用相当清楚，如果性能真的很关键，我认为这是唯一的方法:)

嗯，(0, 0, 3) 是从哪里来的，最后一个轴应该移到中间，不是吗？

@bereal：先增加轴1的值，然后增加轴2，然后增加轴0。哪些值出现在哪个轴上不受axes参数的影响；它只确定值增加的顺序。

@WolframH 在 OP 的成对示例中，轴的顺序发生了变化。

据我了解，4 轴（最后一个，索引为 2）应该移动到结果的中间，因为 (1, 2, 0) 将它放在那里。所以应该有(0, 3, 0)，而不是(3, 0, 0)。

【参考方案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 个维度，它似乎可以工作。 ;)

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

【参考方案4】：

您是否尝试过计时，看看需要多长时间？你所拥有的不应该比没有重新排序慢很多。

您可以尝试修改必须使用的就地拼接分配。

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

您还可以通过将 tpaxes 设为函数的局部变量而不是全局变量（查找时间较慢）来获得加速

否则我的建议是以某种方式编写您自己的产品功能..

【参考方案5】：

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

【讨论】：

这适用于二维情况，但我希望能够指定轴的反转方式。我已更新问题以更清楚地反映这一点。

@GiovanniFunchal 那么我认为你的方法很好：tuple(i[y] for y in tpaxes)

【参考方案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)])

结果：

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