Конкатенация массивов numpy при условии

Question

Мне было интересно, был ли какой-нибудь тупой способ объединить 2 массива различных len (a и b) вдоль оси, основанной на логическом одном из (len a + b)

Я уже могусделать этот массив с пониманием списка Python (в соответствии с логическим значением) или с созданием нескольких новых массивов (2).

Простой пример:

import numpy as np

a = np.arange(3)
b = np.arange(3, 8)
a_b = np.arange(8) < 3

def np_or_not(a_b, a, b, axis=-1):
    axis = axis // len(a_b)
    aligned_a = np.take(a, np.cumsum(a_b) - 1, axis=axis)
    aligned_b = np.take(b, np.cumsum(~a_b) - 1, axis=axis)
    return np.where(a_b, aligned_a, aligned_b)

assert (np_or_not(a_b, a, b) == np.arange(8)).all()

ND case:

Инициализация

axis = 2
a = np.random.rand(10, 20, 4)
b = np.random.rand(10, 20, 9)
a_b = np.array([True] * a.shape[axis] + [False] * b.shape[axis])
np.random.shuffle(a_b)

Проверки (форматы)

assert len(a.shape) == len(b.shape)
assert all(a.shape[ax] == b.shape[ax] for ax in range(len(a.shape)) if ax != axis)
assert len(a_b) == (a.shape[axis] + b.shape[axis])
assert sum(a_b) == a.shape[axis]
assert sum(~a_b) == b.shape[axis]

Проверяет результаты

result = np_or_not(a_b, a, b, axis=-1)

assert result.shape == tuple(l if ax != axis else len(a_b) for ax, l in enumerate(a.shape))

a_b_indexes = [0, 0]
for index, truth in enumerate(a_b):
    assert (result[..., index] == (a if truth else b)[..., a_b_indexes[1 - truth]]).all(), index
    a_b_indexes[1 - truth] += 1

РЕДАКТИРОВАТЬ : Спасибо за ответ, я заменил укладкупутем конкатенации и предоставил пример прецедента ND.