Search code examples
pythonarraysnumpynumpy-ndarraytensor-indexing

Replace values in a 3d numpy array


There are those two numpy arrays:

a = np.array([
    [
        [1,2,3,0,0],
        [4,5,6,0,0],
        [7,8,9,0,0]
    ],
    [
        [1,3,5,0,0],
        [2,4,6,0,0],
        [1,1,1,0,0]
    ]
])

b = np.array([
    [
        [1,2],
        [2,3],
        [3,4]
    ],
    [
        [4,1],
        [5,2],
        [6,3]
    ]
])

with shapes:

"a" shape: (2, 3, 5), "b" shape: (2, 3, 2)

I want to replace the last two elements from array a with those from array b, e.g.

c = np.array([
    [
        [1,2,3,1,2],
        [4,5,6,2,3],
        [7,8,9,3,4]
    ],
    [
        [1,3,5,4,1],
        [2,4,6,5,2],
        [1,1,1,6,3]
    ]
])

However, np.hstack((a[:,:,:-2], b)) throws a Value Error:

all the input array dimensions except for the concatenation axis must match exactly

and in general doesn't look like it's the correct function to use. Append doesn't work either.

Is there a method in numpy that can do that or do I need to iterate over the arrays with a for loop and manipulate them manually?


Solution

  • Non-overwriting method:

    • a[:,:,-2:] fetches the zeros at the end; use a[:,:,:3].

    • According to the documentation, np.hstack(x) is equivalent to np.concatenate(x, axis=1). Since you want to join the matrices on their innermost rows, you should use axis=2.

    Code:

    >>> np.concatenate((a[:,:,:3], b), axis=2)
    array([[[1, 2, 3, 1, 2],
            [4, 5, 6, 2, 3],
            [7, 8, 9, 3, 4]],
    
           [[1, 3, 5, 4, 1],
            [2, 4, 6, 5, 2],
            [1, 1, 1, 6, 3]]])