CODEWITHSUNDEEP
pythondictionaryargumentskeyword-argument
Esther Y
Esther Y

Reputation: 33

With a dictionary of a variable number of keys, how to obtain a cartesian product of its values?

I have a list of dimensions which is expandable, 

sex = ['male', 'female']
yearold = ['<20', '21-30', '31-40', '41-50', '51-60', '>60']
time = ['9am', '10am', '11am']

How do I create a cartesian product and output as a csv file (from a list of lists) with the name of the list as the header? So far I have tried to convert it to a dictionary:

from collections import defaultdict
d = defaultdict(list)

d['sex'].extend(sex)
d['yearold'].extend(yearold)
d['time'].extend(time)

def cartesian(*arg):

    product=[]

    for element in itertools.product(*arg):
        product.append(element)

    titlerow=[]

    for key, value in d.iteritems() :        
        titlerow.append(key)

    mylist.insert(0,titlerow)

cartesian(sex,yearold,time)

It seems a bit messy to have both a list and a dictionary just for the header.. I understand that there is certain use for **kwargs but am unsure of how to incorporate it in!

Upvotes: 2

Views: 180

Answers (3)

salparadise
salparadise

Reputation: 5805

if you care about the order of the rows you can use the values into an OrderedDict:

from collections import OrderedDict
import itertools,csv

sex = ['male', 'female']
yearold = ['<20', '21-30', '31-40', '41-50', '51-60', '>60']
time = ['9am', '10am', '11am']

d = OrderedDict([('sex',sex),('yearold',yearold),('time', time)])

with open('output.csv','wb') as f:
    w = csv.writer(f)
    w.writerow([k for k in d.keys()])
    for e in itertools.product(*[v for v in d.values()]):
        w.writerow([str(i) for i in e])

Output:

sex,yearold,time
male,<20,9am
male,<20,10am
male,<20,11am
male,21-30,9am
male,21-30,10am
male,21-30,11am
male,31-40,9am
male,31-40,10am
male,31-40,11am
male,41-50,9am
male,41-50,10am
male,41-50,11am
male,51-60,9am
male,51-60,10am
male,51-60,11am
male,>60,9am
male,>60,10am
male,>60,11am
female,<20,9am
female,<20,10am
female,<20,11am
female,21-30,9am
female,21-30,10am
female,21-30,11am
female,31-40,9am
female,31-40,10am
female,31-40,11am
female,41-50,9am
female,41-50,10am
female,41-50,11am
female,51-60,9am
female,51-60,10am
female,51-60,11am
female,>60,9am
female,>60,10am
female,>60,11am

Upvotes: 1

nebffa
nebffa

Reputation: 1549

You want to output a .csv file so use the csv module.

def cartesian(**kwargs):
    with open('data.csv', 'w') as f:
        writer = csv.writer(f)
        writer.writerow(kwargs.keys())

        for element in itertools.product(kwargs.values()):
            writer.writerow(element)

With kwargs.keys() you get 'sex', 'time', etc.

With kwargs.values() you get the values in the dictionary.

Upvotes: 0

HennyH
HennyH

Reputation: 7944

Close, you are right in that you can use **kwargs to take in a dictionary of argument name and value pairs. For example:

def cartesian(**columnsByName):
    print(",".join(columnsByName.keys()))
    print("\n".join(",".join(row) for row in product(*columnsByName.values())))

cartesian(sex=sex,yearold=yearold,time=time)

You can then print out the headers then seperately create each row by iterating over the result of product(*). The output will look like this:

yearold,time,sex
<20,9am,male
<20,9am,female
<20,10am,male
<20,10am,female
<20,11am,male
<20,11am,female
21-30,9am,male
...

Upvotes: 0

Related Questions