python - Python中的R expand.grid() 函数

Question

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

(EDIT) 下面是这个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]}

EDIT 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 * 不会返回任何错误消息)

Answer 1

只需使用列表推导:

>>> [(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 开始的。