CODEWITHSUNDEEP
pythonarraysnumpymatrix
Alex
Alex

Reputation: 73

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

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

Upvotes: 1

Views: 47

Answers (1)

akuiper
akuiper

Reputation: 214927

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

Upvotes: 1

Related Questions