Numpy: how to map an element to an array

Question

Suppose I have a 2D numpy array, say 5*3. Now I would like to map each element i in it to a new array [i, i*i], so the resulting array is 5*3*2.

What is the most efficient (and elegant) way to achieve this purpose?

A naive solution using for:

a = np.arange(15).reshape(5, 3)
r = []
for row in a:
   _row = []
   for i in row:
      _row.append([i, i*i])
   r.append(_row)
return np.array(r)

yatu · Accepted Answer · 2019-07-08 13:30:34Z

2

You could use np.dstack to stack both arrays depth wise:

np.dstack([a, a**2])

a = np.arange(15).reshape(5, 3)

array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14]])

np.dstack([a, a**2])

array([[[  0,   0],
        [  1,   1],
        [  2,   4]],

       [[  3,   9],
        [  4,  16],
        [  5,  25]],
 ...

answered Jul 8, 2019 at 13:30

yatu

88.6k12 gold badges93 silver badges148 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

chenzhongpu Over a year ago

what if the desired map is a complicated function? For example, def func(i) returns [x, y], while this func is not trival like [i, i**2].

yatu Over a year ago

I think it would be more appropiate to ask this in a new question @chenzhongpu , note that this was not included in the question, and the corresponding answer would be quite different. Feel free to ping me here when you do so!

Collectives™ on Stack Overflow

Numpy: how to map an element to an array

1 Answer 1

2 Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Related