Отображение данных между массивами различной формы

Question

У меня есть массив данных с формой (256,256,3), где каждое значение на оси (0,1) имеет значение [1,0,0], [0,1,0],[0,0,1] или [1,1,0].

Я хочу преобразовать это в массив формы (256,256,4) где значения отображаются следующим образом:

[1,0,0] -> [1,0,0,0]
[0,1,0] -> [0,1,0,0]
[0,0,1] -> [0,0,1,0]
[1,1,0] -> [0,0,0,1]

Я пытался использовать al oop следующим образом, но это не работает.

for element in original_array:
    if element == [1, 0, 0]:
        new_array[:,:,0] = 1
    elif element == [0, 1, 0]:
        new_aray[:,:,1] = 1
    elif element == [0, 0, 1]:
        new_array[:,:,2] = 1
    elif element == [1, 1, 0]:
        new_array[:,:,3] = 1

Я также пробовал подобный код но я не уверен, как правильно его реализовать:

new_array[:,:,0][original_array == [1,0,0]] = 1
new_array[:,:,1][original_array == [0,1,0]] = 1
new_array[:,:,2][original_array == [0,0,1]] = 1
new_array[:,:,3][original_array == [1,1,0]] = 1

Как правильно реализовать эти два метода и какие альтернативные методы я могу использовать для выполнения sh этой задачи?

Answer 1

Давайте воспроизведем проблему:

a = numpy.array([[1,0,0], [0,1,0], [0,0,1], [1,1,0]])
b = a[numpy.random.randint(0,4,256*256)].reshape(256,256,3)

Мы можем вставить ваш исходный массив в массив нулей с размером 4 на последней оси:

c = numpy.zeros((b.shape[0], b.shape[1], 4))
c[:,:,:3] = b

Затем ищем случаи где у нас есть 1 с индексами 0 и 1, и установить четвертый столбец на 1, где True.

c[:,:,3] = b[:,:,:2].all(axis=-1)

При этом используется эквивалентность между True и 1 (правая часть - логический массив, но мы присвойте его массиву с плавающей точкой, значение которого равно 1).