Interleaving last axis of n-d numpy arrays

Question

Assume the following arrays are given:

a = array([[[4, 7, 1],
        [3, 8, 1],
        [0, 3, 7],
        [1, 5, 5],
        [1, 7, 6]],

       [[8, 1, 0],
        [2, 1, 4],
        [1, 7, 2],
        [6, 2, 0],
        [3, 2, 6]],

       [[4, 3, 1],
        [0, 3, 9],
        [1, 6, 2],
        [6, 5, 7],
        [1, 3, 2]],

       [[7, 3, 7],
        [8, 1, 3],
        [9, 9, 4],
        [2, 6, 8],
        [4, 4, 5]],

       [[4, 8, 1],
        [2, 0, 3],
        [4, 8, 9],
        [3, 7, 2],
        [7, 3, 1]]])

b = array([[[8, 3, 2],
        [1, 7, 9],
        [0, 0, 6],
        [1, 0, 1],
        [6, 0, 2]],

       [[1, 0, 0],
        [8, 6, 9],
        [8, 2, 2],
        [6, 6, 5],
        [2, 0, 8]],

       [[6, 2, 4],
        [4, 5, 4],
        [4, 7, 4],
        [3, 7, 8],
        [0, 9, 5]],

       [[7, 9, 0],
        [5, 4, 9],
        [7, 5, 9],
        [0, 0, 5],
        [4, 7, 2]],

       [[0, 2, 0],
        [4, 7, 7],
        [6, 7, 6],
        [2, 5, 3],
        [6, 0, 5]]])

How can i interleave the row to get shape to be (5, 5, 6) and the value on the last axis to be

4, 8, 7, 3, 1, 2

This is repeated for the rest of the row. I have tried examples from Interweaving two numpy arrays but it doesn't work since the reshape cause the shape mismatch. A numpy solution would be greatly appreciated.

Doesn't the accepted answer to the linked question do exactly what you want? `c = np.empty((5, 5, 6), dtype=a.dtype); c[...,0::2] = a; c[...,1::2] = b; c[0,0]`? — Reti43, Jan 12 '21 at 01:12
I'd try reshaping them to (5,5,3,1), concatenate on last axis to get (5,5,3,2). Finally reshape to (5,5,6)` — hpaulj, Jan 12 '21 at 01:33
@Reti43 that's a good point, is there a way to do it with operations such as dstack etc? I want to avoid having to hardcode the indexes — AlwaysNull, Jan 12 '21 at 01:45
You don't have to hardcode it, I just did it for convenience. `s = a.shape; new_shape = (s[0], s[1], 2*s[2])`. Or if you want to keep it even more generic, `new_shape = s[:-1] + (2*s[-1],)`. — Reti43, Jan 12 '21 at 02:05

score 2 · Accepted Answer · answered Jan 12 '21 at 02:24

You can stack the two arrays along a new axis first (i.e. the 4th axis here), and then flatten the last two dimensions which will interleave elements on the 3rd axis from the original arrays:

np.stack((a, b), axis=-1).reshape(a.shape[:-1] + (-1,))

#[[[4 8 7 3 1 2]
#  [3 1 8 7 1 9]
#  [0 0 3 0 7 6]
#  [1 1 5 0 5 1]
#  [1 6 7 0 6 2]]
# ...

Interleaving last axis of n-d numpy arrays

1 Answers1