Make new 2D array in third dimension for every element in 1D array, pythonically

Question

Say you have two 1D arrays with the same size a=[1,8,4,5,9] b=[1,8,4,5,9] and then for every ith element in a and b, can you make a new array such that,

where H is a matrix of matrices and H(i) is stacked in a third dimension?

I have tried using numpys np.dstack however it seems like it treats every new element inputted separately.
Doing this with a for loop would be trivial, however I believe that they are slow in python and thus wish to do this in a pythonic way, with numpy if possible.

so then for H[0] we would have [[1,16,16],[1,7,1],[1,4,2]]
and similarly for H[1] we would have [[64,16,16],[64,56,64],[8,32,16]]

Could you clarify whether you're looking for a solution that can use numpy, or do you want no imports? — Josh Karpel
– Josh Karpel, Commented Aug 31, 2017 at 21:04

akuiper · Accepted Answer · 2017-08-31 21:08:57Z

1

Use column_stack to stack the calculated results and then reshape:

a=np.array([1,8,4,5,9])
b=np.array([1,8,4,5,9])

np.column_stack((
    a ** 2, 2 * a, 2 * a,
    b * a,  7 * a, b * a,
    b,      4 * b, 2 * b  
)).reshape(-1,3,3)

Out[468]:
array([[[ 1,  2,  2],
        [ 1,  7,  1],
        [ 1,  4,  2]],

       [[64, 16, 16],
        [64, 56, 64],
        [ 8, 32, 16]],

       [[16,  8,  8],
        [16, 28, 16],
        [ 4, 16,  8]],

       [[25, 10, 10],
        [25, 35, 25],
        [ 5, 20, 10]],

       [[81, 18, 18],
        [81, 63, 81],
        [ 9, 36, 18]]])

answered Aug 31, 2017 at 21:08

akuiper

216k33 gold badges362 silver badges379 bronze badges

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

1 Comment

Alex Over a year ago

This is rather beautiful. And its pretty much a one liner, rather than having a for loop style iteration.

Collectives™ on Stack Overflow

Make new 2D array in third dimension for every element in 1D array, pythonically

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related