Numpy 2d array - выбрать несколько элементов без цикла for

Question

Я поддерживаю некоторый код, и я сталкиваюсь с чем-то вроде:

    travel_time_vec = np.zeros(...)
    for v in some_indexes: # some_indexes is a list of row indexes
        traveltimes = traveltime_2d_array[v, list_of_column_indexes]
        best_index = np.argmin(traveltimes)
        travel_time_vec[v] = traveltimes[best_index]

Я хотел бы отбросить цикл for и выполнить все операции ниже, но наивно спрашивая traveltime_2d_array[some_indexes, list_of_column_indexes] результаты в:

{IndexError} несоответствие форм: индексные массивы не могут быть переданы вместе с формами (4,) (8,)

Answer 1

Понял - мне нужно пройти some_indexes в виде списка списков, так что NumPy транслирует каждый столбцы в list_of_column_indexes, Итак, это:

travel_time_vec = np.zeros(...)
# newaxis below tranforms [1, 2, 3] to [[1], [2], [3]]
traveltimes = traveltime_2d_array[np.array(some_indexes)[:, np.newaxis], 
                                  list_of_column_indexes]
# get the index of the min time on each row
best_index = np.argmin(traveltimes, axis=1)
travel_time_vec[some_indexes] = traveltimes[:, best_index]

работает как положено, и больше не зацикливается