Как включить np.argmax в трансляцию, чтобы заменить данный код цикла for (если это возможно)

Question

Я играл с Numpy, чтобы ускорить код, где смогу.Это действительно красиво быстро.Однако иногда это требует умного мышления.Я полагаю, что практика станет совершенной.

Вместо того, чтобы объяснять мою проблему, вот что я пытаюсь заменить:

# We use U and A to compute V

U = np.array([[1,2],
             [4,3],
             [5,6],
             [7,8]])

V = np.zeros(U.shape)

A = np.array([[1,3],
              [3,4]])

# The for loop to be replaced

for t in range(len(U)):
    V[t] = np.argmax( U[t]*A.T ,axis = 1)

Моя попытка:

V = np.argmax(U[:,np.newaxis]*A.T,axis=1)

# U[:,np.newaxis]*A.T 

Не бойся, я знаю, что не так в моей версии Numpy.Закомментированный код дает правильное промежуточное звено, однако мне не удается правильно включить часть np.argmax, как в моем коде цикла for.Я не думаю, что можно, возможно.Пожалуйста, помогите мне, если это возможно.Я очень ценю это.

Answer 1

Вы на правильном пути, просто неправильная ось.Также вам не нужно транспонировать A:

(A * U[:,np.newaxis]).argmax(-1) # or equivalently axis 2

array([[1, 1],
       [1, 0],
       [1, 1],
       [1, 1]], dtype=int64)