best way to create numpy array from FOR loop

Question

Is there a better way to create a multidimensional array in numpy using a FOR loop, rather than creating a list? This is the only method I could come up with:

import numpy as np

a = []
for x in range(1,6):
    for y in range(1,6):
        a.append([x,y])
a = np.array(a)
print(f'Type(a) = {type(a)}.  a = {a}')

EDIT: I tried doing something like this:

a = np.array([range(1,6),range(1,6)])
a.shape = (5,2)
print(f'Type(a) = {type(a)}.  a = {a}')

however, the output is not the same. I'm sure I'm missing something basic.

Do you understand what that 2nd try did? Even it didn't produce what you want, it's important to understand that expression. — hpaulj, Jan 05 '18 at 18:21

aasoo · Answer 1 · 2020-09-12T21:56:40.460

3

You could preallocate the array before assigning the respective values:

a = np.empty(shape=(25, 2), dtype=int)
for x in range(1, 6):
    for y in range(1, 6):
        index = (x-1)*5+(y-1)
        a[index] = x, y

edited Sep 12 '20 at 21:56

answered Sep 12 '20 at 21:44

aasoo

100
1
8

score 1 · Answer 2 · answered Jan 05 '18 at 15:29

1

Did you had a look at numpy.ndindex? This could do the trick:

a = np.ndindex(6,6)

You could have some more information on Is there a Python equivalent of range(n) for multidimensional ranges?

answered Jan 05 '18 at 15:29

Tobias Braun

41
3

koPytok · Answer 3 · 2021-01-13T10:12:44.273

1

You can replace double for-loop with itertools.product.

from itertools import product 
import numpy as np

np.array(list(product(range(1,6), range(1,6))))

For me creating array from list looks natural here. Don't know how to skip them in that case.

edited Jan 13 '21 at 10:12

answered Jan 05 '18 at 16:00

koPytok

3,453
1
14
29

Yerbol Sapar · Answer 4 · 2022-09-27T06:05:27.303

Sometimes it is difficult to predict the total element number and shape at stage of selecting array elements due to some if statement inside a loop.

In this case, put all selected elements in flat array:

a = np.empty((0), int)
for x in range(1,6):  # x-coordinate
    for y in range(1,6):  # y-coordinate
        if x!=y:  # `if` statement
            a = np.append(a, [x, y])

Then, given the lengths of one array dimension (in our case there are 2 coordinates) one can use -1 for the unknown dimension:

a.shape = (-1, 2)
a
array([[1, 2],
       [1, 3],
       [1, 4],
       [1, 5],
       [2, 1],
       [2, 3],
       [2, 4],
       [2, 5],
       [3, 1],
       [3, 2],
       [3, 4],
       [3, 5],
       [4, 1],
       [4, 2],
       [4, 3],
       [4, 5],
       [5, 1],
       [5, 2],
       [5, 3],
       [5, 4]])

best way to create numpy array from FOR loop

4 Answers4