python - Numpy 2d 数组 - 在没有 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]

按预期工作，不再循环