Search code examples
pythonarraysnumpytile

2D of numpy arange

I want to build two 2d arrays

a = [[0, 0, 0, 0, 0, 0],
     [1, 1, 1, 1, 1, 1],
     [2, 2, 2, 2, 2, 2],
     [3, 3, 3, 3, 3, 3],
     [4, 4, 4, 4, 4, 4],
     [5, 5, 5, 5, 5, 5]]

b = [[0, 1, 2, 3, 4, 5],
     [0, 1, 2, 3, 4, 5],
     [0, 1, 2, 3, 4, 5],
     [0, 1, 2, 3, 4, 5],
     [0, 1, 2, 3, 4, 5],
     [0, 1, 2, 3, 4, 5]]

But i cannot use for loops. i know i can get an array using np.arange(5), but not sure how to turn that into the 2 2D arrays shown above. Any help would be greatly appreciated

Solution

  • You can use numpy.mgrid or numpy.meshgrid():

    np.mgrid[0:6, 0:6]
# array([[[0, 0, 0, 0, 0, 0],
#         [1, 1, 1, 1, 1, 1],
#         [2, 2, 2, 2, 2, 2],
#         [3, 3, 3, 3, 3, 3],
#         [4, 4, 4, 4, 4, 4],
#         [5, 5, 5, 5, 5, 5]],
# 
#        [[0, 1, 2, 3, 4, 5],
#         [0, 1, 2, 3, 4, 5],
#         [0, 1, 2, 3, 4, 5],
#         [0, 1, 2, 3, 4, 5],
#         [0, 1, 2, 3, 4, 5],
#         [0, 1, 2, 3, 4, 5]]])

np.meshgrid(np.arange(6), np.arange(6))
# [array([[0, 1, 2, 3, 4, 5],
#         [0, 1, 2, 3, 4, 5],
#         [0, 1, 2, 3, 4, 5],
#         [0, 1, 2, 3, 4, 5],
#         [0, 1, 2, 3, 4, 5],
#         [0, 1, 2, 3, 4, 5]]),
#  array([[0, 0, 0, 0, 0, 0],
#         [1, 1, 1, 1, 1, 1],
#         [2, 2, 2, 2, 2, 2],
#         [3, 3, 3, 3, 3, 3],
#         [4, 4, 4, 4, 4, 4],
#         [5, 5, 5, 5, 5, 5]])]

    and simply unpack the result

    a, b = np.mgrid[0:6, 0:6]