Suggestions for storing outputs of a nested forloop?

Viewed 44

I am trying to write a function in python that gives all possible combinations of three variable inputs that can be of size 1-4. I have written a nested forloop that I believe gives all possible combinations, but I am struggling to store the output in a single 1D array. I don't know if I have to store it as a 3D array and compress it maybe? Here is the code I have:

import numpy as np
def test(x,y,z):
    len1 = len(x)
    len2 = len(y)
    len3 = len(z)
    lentot = len1*len2*len3
    codons = np.empty((1,lentot))
    for i in range(len1):
        for j in range(len2):
            for k in range(len3):
                codons[] = np.array([x[i],y[j],z[k]])
    return codons

Basically, I a cannot figure out what to put in the bracket on the second to last line to get my output to store as a 1D array. I don't even know if it is possible. I tried using itertools.product to perform this for me, but the output is stored as a single element, not an array (each line being its own element). For my application it is important that I can pass this output through another function, so I need it to be an array of strings.

3 Answers

You can append them to a list and then convert it to an array at the end.

def test(x, y, z):
    codon_list = []
    for i in x:
        for j in y:
            for k in z:
                codon_list.append([i, j, k])
    codons = np.array(codon_list)

Using the combinations function from itertools itertools.combinations([x, y, z], 3) seems to be the easiest way.

You can use list comprehension to make your life easier:

def test(x,y,z):
    codons = np.array([[i,j,k] for i in x for j in y for k in z])
    return codons
Related