Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Create a 2-D numpy array with list comprehension

Tags:

python

numpy

I need to create a 2-D numpy array using only list comprehension, but it has to follow the following format:

[[1, 2, 3],
 [2, 3, 4],
 [3, 4, 5],
 [4, 5, 6],
 [5, 6, 7]]]

So far, all I've managed to figure out is:

two_d_array = np.array([[x+1 for x in range(3)] for y in range(5)])

Giving:

array([[1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3],
       [1, 2, 3]])

Just not very sure how to change the incrementation. Any help would be appreciated, thanks!

EDIT: Accidentally left out [3, 4, 5] in example. Included it now.

like image 229
ambrrrgris Avatar asked Feb 24 '26 21:02

ambrrrgris


2 Answers

Here's a quick one-liner that will do the job:

np.array([np.arange(i, i+3) for i in range(1, 6)])

Where 3 is the number of columns, or elements in each array, and 6 is the number of iterations to perform - or in this case, the number of arrays to create; which is why there are 5 arrays in the output.

Output:

array([[1, 2, 3],
       [2, 3, 4],
       [3, 4, 5],
       [4, 5, 6],
       [5, 6, 7]])
like image 173
S3DEV Avatar answered Feb 26 '26 09:02

S3DEV


Change the code, something like this can work:

two_d_array = np.array([[(y*3)+x+1 for x in range(3)] for y in range(5)])
>>> [[1,2,3],[4,5,6],...]
two_d_array = np.array([[y+x+1 for x in range(3)] for y in range(5)])
>>> [[1,2,3],[2,3,4],...]
like image 31
Uriya Harpeness Avatar answered Feb 26 '26 10:02

Uriya Harpeness