Numpy сравнить два 3D-массива и найти идентичные массивы

Question

У меня есть 2 массива,

mat = np.array(
    [
        [[0,0],[0,1],[0,2]],
        [[1,0],[1,1],[1,2]],
        [[2,0],[2,1],[2,2]]
    ]
)
mat2 = np.array(
    [
        [[0,1],[0,1],[0,2]],
        [[1,0],[1,1],[1,2]],
        [[2,0],[2,2],[2,2]]
    ]
)

Я хотел бы найти все массивы на оси 2, которые одинаковы как в mat, так и mat2, и превратить их в массив нулей. Иными словами, если бы каждый массив mat и mat2 был изображением RGB, я хотел бы найти пиксели с одинаковыми значениями R, G, B и вернуть изображение с такими же пикселями, установленными как (0,0,0 ) и остальные пиксели не повреждены.

Таким образом, вывод, который я пытаюсь получить с помощью вышеуказанных массивов, будет следующим:

[
    [[0, 1], [0, 0], [0, 0]],
    [[0, 0], [0, 0], [0, 0]],
    [[0, 0], [2, 2], [0, 0]]
]

Я пробовал следующий код:

operated = np.where((mat2-mat==0).all(axis=2), np.array([0,0]), mat2)

Однако он сказал бы:

ValueError: operands could not be broadcast together with shapes (3,3) (2,) (3,3,2) 

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

Answer 1

Вы можете использовать np.where для замены значений в массиве в соответствии со следующим условием:

np.where((mat == mat2).all(axis=2,  keepdims=True), [0,0], mat2)

array([[[0, 1],
        [0, 0],
        [0, 0]],

       [[0, 0],
        [0, 0],
        [0, 0]],

       [[0, 0],
        [2, 2],
        [0, 0]]])

Где (mat == mat2).all(axis=2, keepdims=True):

array([[[False],
        [ True],
        [ True]],

       [[ True],
        [ True],
        [ True]],

       [[ True],
        [False],
        [ True]]])

Возвращает маску тех же размеров, что и mat, которую можно использовать в качестве условия для np.where. Затем вам просто нужно указать, заменять ли эти значения на [0,0] или mat2 в зависимости от результата.