Choose axis for assignment into NumPy array programmatically

I have a 3D numpy array.

I can modify arbitrary element using simple indexing


Sometimes I need to change element(s) at a given index and it is not predetermined at which axis the index refers. I would like to catch the two following cases with a change in variable

D[:,:10,1]=list(range(10)) ->axis 1
D[:10,:,1]=list(range(10)) ->axis 0

Something like:


1 answer

  • answered 2018-11-08 18:17 Welcome to Stack Overflow

    I'd use an indexing tuple with slice objects prepared by the helper object np.s_. If axis is 0 or 1, the following has the effect of assigning list(range(10)) to either D[:10, :, 1] or D[:, :10, 1].

    idx = [np.s_[:], np.s_[:], 1] 
    idx[axis] = np.s_[:10]
    D[tuple(idx)] = list(range(10))