CODEWITHSUNDEEP
pythonnumpy
ambrrrgris
ambrrrgris

Reputation: 157

Create a 2-D numpy array with list comprehension

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.

Upvotes: 3

Views: 202

Answers (3)

Mad Physicist
Mad Physicist

Reputation: 114350

You've got a couple of good comprehension answers, so here are a couple of numpy solutions.

Simple addition:

np.arange(1, 6)[:, None] + np.arange(3)

Crazy stride tricks:

base = np.arange(1, 8)
np.lib.stride_tricks.as_strided(base, shape=(5, 3), strides=base.strides * 2).copy()

Reshaped cumulative sum:

base = np.ones(15)
base[3::3] = -1
np.cumsum(base).reshape(5, 3)

Upvotes: 1

s3dev
s3dev

Reputation: 9701

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]])

Upvotes: 1

Uriya Harpeness
Uriya Harpeness

Reputation: 628

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],...]

Upvotes: 1

Related Questions