Python 中的 R expand.grid() 函数

Question

是否有类似于 R 中的 expand.grid() 函数的 Python 函数？提前致谢。

（编辑）下面是这个 R 函数的描述和一个例子。

 Create a Data Frame from All Combinations of Factors

Description:

     Create a data frame from all combinations of the supplied vectors
     or factors.

> x <- 1:3
> y <- 1:3
> expand.grid(x,y)
  Var1 Var2
1    1    1
2    2    1
3    3    1
4    1    2
5    2    2
6    3    2
7    1    3
8    2    3
9    3    3

(EDIT2) 以下是 rpy 包的示例。我想获得相同的输出对象但不使用 R ：

 >>> from rpy import *
>>> a = [1,2,3]
>>> b = [5,7,9]
>>> r.assign("a",a)
[1, 2, 3]
>>> r.assign("b",b)
[5, 7, 9]
>>> r("expand.grid(a,b)")
{'Var1': [1, 2, 3, 1, 2, 3, 1, 2, 3], 'Var2': [5, 5, 5, 7, 7, 7, 9, 9, 9]}

编辑 02/09/2012： 我真的迷失了 Python。 Lev Levitsky 在他的回答中给出的代码对我不起作用：

 >>> a = [1,2,3]
>>> b = [5,7,9]
>>> expandgrid(a, b)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 2, in expandgrid
NameError: global name 'itertools' is not defined

但是似乎安装了 itertools 模块（键入 from itertools import * 不返回任何错误消息）

原文由 Stéphane Laurent 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

这是一个示例，它提供类似于您需要的输出：

 import itertools
def expandgrid(*itrs):
   product = list(itertools.product(*itrs))
   return {'Var{}'.format(i+1):[x[i] for x in product] for i in range(len(itrs))}

>>> a = [1,2,3]
>>> b = [5,7,9]
>>> expandgrid(a, b)
{'Var1': [1, 1, 1, 2, 2, 2, 3, 3, 3], 'Var2': [5, 7, 9, 5, 7, 9, 5, 7, 9]}

不同之处在于 itertools.product 最右边的元素在每次迭代中都会前进。如果它很重要，您可以通过巧妙地对 product 列表进行排序来调整该功能。

---

编辑（S. Laurent）

与 R 相同：

 def expandgrid(*itrs): # https://stackoverflow.com/a/12131385/1100107
    """
    Cartesian product. Reversion is for compatibility with R.

    """
    product = list(itertools.product(*reversed(itrs)))
    return [[x[i] for x in product] for i in range(len(itrs))][::-1]

原文由 Lev Levitsky 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 2

只需使用列表理解：

 >>> [(x, y) for x in range(5) for y in range(5)]

[(0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (1, 0), (1, 1), (1, 2), (1, 3), (1, 4), (2, 0), (2, 1), (2, 2), (2, 3), (2, 4), (3, 0), (3, 1), (3, 2), (3, 3), (3, 4), (4, 0), (4, 1), (4, 2), (4, 3), (4, 4)]

如果需要，转换为 numpy 数组：

 >>> import numpy as np
>>> x = np.array([(x, y) for x in range(5) for y in range(5)])
>>> x.shape
(25, 2)

我已经测试了高达 10000 x 10000 并且 python 的性能与 R 中的 expand.grid 相当。在理解中使用元组 (x, y) 比使用列表 [x, y] 快大约 40%。

或者…

使用 np.meshgrid 大约快 3 倍，而且内存占用更少。

 %timeit np.array(np.meshgrid(range(10000), range(10000))).reshape(2, 100000000).T
1 loops, best of 3: 736 ms per loop

在 R 中：

 > system.time(expand.grid(1:10000, 1:10000))
   user  system elapsed
  1.991   0.416   2.424

请记住，R 具有从 1 开始的数组，而 Python 是从 0 开始的。

原文由 Thomas Browne 发布，翻译遵循 CC BY-SA 4.0 许可协议