Numpy 2d 数组 - 选择没有 for 循环的多个元素 [关闭]答案

【问题标题】：Numpy 2d array - select multiple elements without a for loop [closed]Numpy 2d 数组 - 选择没有 for 循环的多个元素 [关闭]
【发布时间】：2018-11-07 09:33:34
【问题描述】：

我维护了一些代码，但遇到了类似的情况：

    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,) 一起广播

【问题讨论】：

不知道你在这里做什么。你的问题听起来像“这是一个写得很糟糕的函数，现在为我重构它”。你能提供一个minimal reproducible example 描述你在做什么吗？
traveltimes = traveltime_2d_array[v, list_of_column_indexes] 将traveltime_2d_array 的元素数组返回给我，这些元素位于v 的行list_of_column_indexes 上。我想矢量化 for 循环 - 也就是在some_indexes 上为每一行返回一个二维数组
另见stackoverflow.com/q/22927181/281545

标签： python numpy numpy-indexing

【解决方案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]

按预期工作，不再循环

【讨论】：

没错，一个(n,1)数组可以用一个(m,)数组广播来选择(n,m)个元素。或者 (m,) 可以与 (m,) 一起选择 (m,) 元素（例如沿对角线）。
@hpaulj：如果你愿意，你可以尝试回答:)