Access elements of a Matrix by a list of indices in Python to apply a max(val, 0.5) to each value without a for loop

Question

I know how to access elements in a vector by indices doing:

test = numpy.array([1,2,3,4,5,6])
indices = list([1,3,5])
print(test[indices])

which gives the correct answer : [2 4 6]

But I am trying to do the same thing using a 2D matrix, something like:

currentGrid = numpy.array(  [[0,   0.1],
                             [0.9, 0.9],
                             [0.1, 0.1]])
indices = list([(0,0),(1,1)])
print(currentGrid[indices])

this should display me "[0.0 0.9]" for the value at (0,0) and the one at (1,1) in the matrix. But instead it displays "[ 0.1 0.1]". Also if I try to use 3 indices with :

indices = list([(0,0),(1,1),(0,2)])

I now get the following error:

Traceback (most recent call last):
  File "main.py", line 43, in <module>
    print(currentGrid[indices])
IndexError: too many indices for array

I ultimately need to apply a simple max() operation on all the elements at these indices and need the fastest way to do that for optimization purposes.

What am I doing wrong ? How can I access specific elements in a matrix to do some operation on them in a very efficient way (not using list comprehension nor a loop).

Answer 1

The problem is the arrangement of the indices you're passing to the array. If your array is two-dimensional, your indices must be two lists, one containing the vertical indices and the other one the horizontal ones. For instance:

idx_i, idx_j = zip(*[(0, 0), (1, 1), (0, 2)])
print currentGrid[idx_j, idx_i]
# [0.0, 0.9, 0.1]

Note that the first element when indexing arrays is the last dimension, e.g.: (y, x). I assume you defined yours as (x, y) otherwise you'll get an IndexError