Is there a Python function similar to the expand.grid() function in R ? Thanks in advance.
(EDIT) Below are the description of this R function and an example.
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) Below is an example with the rpy package. I would like to get the same output object but without using 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: I'm really lost with Python. Lev Levitsky's code given in his answer does not work for me:
>>> a = [1,2,3]
>>> b = [5,7,9]
>>> expandgrid(a, b)
Traceback (most recent call last):
File "", line 1, in
File "", line 2, in expandgrid
NameError: global name 'itertools' is not defined
However the itertools module seems to be installed (typing from itertools import *
does not return any error message)
Answer
Here's an example that gives output similar to what you need:
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]}
The difference is related to the fact that in itertools.product
the rightmost element advances on every iteration. You can tweak the function by sorting the product
list smartly if it's important.
No comments:
Post a Comment