Python: найти максимальное значение вдоль одной оси в 3d-матрице, сделать ненулевые значения равными нулю

Question

У меня есть трехмерная матрица, как показано ниже, и я хотел бы взять максимальное значение вдоль оси 1 и сохранить все не максимальные значения равными нулю.

A = np.random.rand(3,3,2)

  [[[0.34444547, 0.50260393],
    [0.93374423, 0.39021899],
    [0.94485653, 0.9264881 ]],

   [[0.95446736, 0.335068  ],
    [0.35971558, 0.11732342],
    [0.72065402, 0.36436023]],

   [[0.56911013, 0.04456443],
    [0.17239996, 0.96278067],
    [0.26004909, 0.06767436]]]

Желаемый результат:

   [[0         , 0         ],
    [0         , 0         ],
    [0.94485653, 0.9264881]],

   [[0.95446736, 0          ],
    [0         , 0          ],
    [0         , 0.36436023]],

   [[0.56911013, 0         ],
    [0         , 0.96278067],
    [0         , 0         ]]])

Я пробовал:

B = np.zeros_like(A)  #return matrix of zero with same shape as A

max_idx = np.argmax(A, axis=1) #index along axis 1 with max value

    array([[2, 0],
           [2, 2],
           [0, 2],
           [0, 1]])

C = np.max(A, axis=1, keepdims = True) #gives a (4,1,2) matrix of max value along axis 1

    array([[[0.95377958, 0.92940525]],
           [[0.94485653, 0.9264881 ]],
           [[0.95446736, 0.36436023]],
           [[0.56911013, 0.96278067]]])

Но я не могу понять, как объединить эти идеи, чтобы получить желаемый результат. Пожалуйста, помогите !!

Answer 1

Вы можете получить 3-мерный индекс ваших максимальных значений из max_idx.Значения в max_idx являются индексом вдоль оси 1 ваших максимальных значений.Есть шесть значений, так как ваши другие оси 3 и 2 (3 x 2 = 6).Вам просто нужно осознать порядок, по которому numpy проходит через них, чтобы получить индекс для каждой из других осей.Сначала вы выполняете итерацию по последним осям:

d0, d1, d2 = A.shape
a0 = [i for i in range(d0) for _ in range(d2)]   # [0, 0, 1, 1, 2, 2]
a1 = max_idx.flatten()                           # [2, 2, 0, 2, 0, 1]
a2 = [k for _ in range(d0) for k in range(d2)]   # [0, 1, 0, 1, 0, 1]

B[a0, a1, a2] = A[a0, a1, a2]

Вывод:

array([[[0.        , 0.        ],
        [0.        , 0.        ],
        [0.94485653, 0.9264881 ]],

       [[0.95446736, 0.        ],
        [0.        , 0.        ],
        [0.        , 0.36436023]],

       [[0.56911013, 0.        ],
        [0.        , 0.96278067],
        [0.        , 0.        ]]])