Multiplying 3 arrays to form a 3D array with entries multiplied

Viewed 71

Say we have three 1D arrays A (m elements), B (n elements) and C (k elements).

Is there any python numpy function that can take the above as input, and output a 3D matrix (dimension: m x n x k) whereby each entry indexed [i, j, k] is of value A[i] * B[j] * C[k] ?

np.outer(A, B) would work for 2D, but I'm looking for a 3D version.

1 Answers

Have a look at numpy broadcasting. Understanding this concepts allows you to tackle most problems regarding multidimensional arrays:

import numpy as np

n, m, k = 13, 17, 19
a = np.random.random(n)
b = np.random.random(m)
c = np.random.random(k)

res = (a[         :, np.newaxis, np.newaxis] *  #   (n, 1, 1)
       b[np.newaxis,          :, np.newaxis] *  # * (1, m, 1)
       c[np.newaxis, np.newaxis,          :])   # * (1, 1, k) 
                                                # = (n, m, k)
print(res.shape)
# (13, 17, 19)

If you are familiar with Einstein Summation Convention einsum might also be interesting for you:

res = np.einsum('n, m, k -> nmk', a, b, c)
Related