У меня есть 3D-массив r (1000 x 10 x 2000)
, построенный следующим образом:
q = np.random.normal(size=(10,2000))
r = np.random.normal(loc=q, size=(1000,10,2000))
Этот массив, r
, можно просмотреть как матрицу 1000 x 10
, повторенную 2000 раз.
Я хотел бы уменьшить этот массив в соответствии со следующим правилом:
- из каждой матрицы выберите только столбец с максимальным значением в первой строке
Столбцы длябыть выбранным можно получить через: np.argmax(r[0], axis=0)
.
Результат должен быть 1000 x 2000
матрица.
Интересно, возможно ли получить что-то подобное без использования цикла for
или списков.
Вот цикл for
, который выполняет вышеуказанную задачу:
x = []
for i, idx in enumerate(np.argmax(r[0], axis=0)):
x.append(r[:,idx,i])
x = np.array(x).T