pythonnumpyarray-broadcasting

Alternative to looping over one numpy axis

I have two numpy arrays a and b such that a.shape[:-1] and b.shape are broadcastable. With this constraint only, I want to calculate an array c according to the following:

c = numpy.empty(numpy.broadcast_shapes(a.shape[:-1],b.shape),a.dtype)
for i in range(a.shape[-1]):
    c[...,i] = a[...,i] * b

The above code certainly works, but I would like to know if there is a more elegant (and idiomatic) way of doing it.

Solution

  • Use np.newaxis with ... to add a new axis after your last axis.

    c = a * b[..., np.newaxis]

    Which is the same as

    c = a * b[np.newaxis, :]

    You don't need to allocate space for c in advance btw.