Numpy 2d array - select multiple elements without a for loop

Question

I maintain some code and I run across something like:

    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]

I would like to drop the for loop and do all the operations below at once - but naively asking for traveltime_2d_array[some_indexes, list_of_column_indexes] results in:

{IndexError}shape mismatch: indexing arrays could not be broadcast together with shapes (4,) (8,)

Not sure what you're upto here. Your question sounds like "this is a badly written function, refactor it for me now". Can you provide a [mcve] describing what you're doing? — cs95, May 28 '18 at 19:11
`traveltimes = traveltime_2d_array[v, list_of_column_indexes]` gives me back an array of the elements of `traveltime_2d_array` that are on row `v` at coulms `list_of_column_indexes`. I want to vectorize the for loop - aka get a 2d array back one for each row on `some_indexes` — Mr_and_Mrs_D, May 28 '18 at 19:15

score 0 · Accepted Answer · answered May 28 '18 at 19:58

0

Got it - I need to pass the some_indexes as a list of lists so numpy broadcasts each one to the columns in list_of_column_indexes. So this:

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]

works as expected, and no more looping

answered May 28 '18 at 19:58

Mr_and_Mrs_D

32,208
39
178
361

That's right, a (n,1) array can broadcast with a (m,) array to select (n,m) elements. Or a (m,) can work with (m,) to select (m,) elements (such as along a diagonal). – hpaulj May 28 '18 at 20:51
@hpaulj: you could attempt an answer if you will :) – Mr_and_Mrs_D May 28 '18 at 22:48

Numpy 2d array - select multiple elements without a for loop

1 Answers1