pythonarraysnumpyarray-broadcasting

Shortest way to broadcast 1d array to specific dimension in NumPy


I often find myself to broadcast 1d arrays into a specific dimension dim_a given a total number of dimension dim_total. What I mean is the following:

import numpy as np

a = np.arange(10)

dim_a = 2
dim_total = 4
shape = tuple([-1 if idx == dim_a else 1 for idx in range(dim_total)]) 
print(a.reshape(shape))


axis = list(range(dim_total))
del axis[dim_a]
print(np.expand_dims(a, axis=axis))

Both work as expected, however the question is whether there is an even shorter way to achieve this for a single array?


Solution

  • Shorter way to get shape:

    shape, shape[dim_a] = [1] * dim_total, -1
    

    Attempt This Online!

    Though I'd prefer it in two lines:

    shape = [1] * dim_total
    shape[dim_a] = -1