Search code examples
pythonpython-3.xnumpymultidimensional-arrayexport-to-csv

create a np.array in a simpler way, and transferring to csv


I am looking for a simpler code to create the following array:

[[1. 1. 1.]
 [2. 1. 1.]
 [3. 1. 1.]
 [4. 1. 1.]
 [1. 2. 1.]
 [2. 2. 1.]
 [3. 2. 1.]
 [4. 2. 1.]
 [1. 3. 1.]
 [2. 3. 1.]
 [3. 3. 1.]
 [4. 3. 1.]
 [1. 1. 2.]
 [2. 1. 2.]
 [3. 1. 2.]
 [4. 1. 2.]
 [1. 2. 2.]
 [2. 2. 2.]
 [3. 2. 2.]
 [4. 2. 2.]
 [1. 3. 2.]
 [2. 3. 2.]
 [3. 3. 2.]
 [4. 3. 2.]]

I have written the following code to generate that, but definitely, it is not a good one.

import numpy as np

t1 = np.array([])
t2 = np.array([])

# creating the i column

for k in range (1,3):
    for j in range(1,4):
        x = np.array(list(range(1, 5)))
        y = np.append(t1, x)
        t1 = y

# creating the j column

for k in range(1,3):
    x = np.array(list(range(1, 4)))
    y = np.repeat(x, 4)
    t2 = np.append(t2, y)

# creating the k column

x = np.array(list(range(1, 3)))
t3 = np.repeat(x, 3*4)
new = np.transpose(np.vstack((t1, t2, t3)))

Also, in the next step I want to export it to a csv file in a way that these numbers goes there in this shape ( I mean my first column in csv file includes: 1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1,2,3 and so on for the next two columns) I tried to use numpy.savetxt() but it seems it only works for 1D and 2D arrays.

Any suggestions?


Solution

  • The itertools.product fits your use-case perfectly:

    from itertools import product
    import pandas as pd
    
    X = sorted(list(product([1,2,3,4], [1,2,3], [1,2])), key=lambda x: (x[2], x[1], x[0]))
    
    pd.DataFrame(X).to_csv('output.csv', header=False, index=False)
    

    Output:

    [[1, 1, 1], [2, 1, 1], [3, 1, 1], [4, 1, 1], [1, 2, 1], [2, 2, 1], [3, 2, 1], [4, 2, 1], [1, 3, 1], [2, 3, 1], [3, 3, 1], [4, 3, 1], [1, 1, 2], [2, 1, 2], [3, 1, 2], [4, 1, 2], [1, 2, 2], [2, 2, 2], [3, 2, 2], [4, 2, 2], [1, 3, 2], [2, 3, 2], [3, 3, 2], [4, 3, 2]]