berkelem berkelem - 4 months ago 8
Python Question

Efficient way to match coordinates in 2d array

I have a 2d array of coordinates and I want to find the index of the entry that matches a given coordinate.

For example, my array could be

A
:

A = [[[1.5, 2.0], [1.0, 2.3], [5.4, 2.3]],
[[3.2, 4.4], [2.0, 3.1], [0.0, 2.3]],
[[1.0, 2.0], [2.3, 3.4], [4.0, 1.1]]]


and the coordinate I want to match is
x = [1.0, 2.0]
. I want to get the index of the coordinate
[1.0, 2.0]
, which would be
(2, 0)
.

Currently I am doing it as follows:

matching_inds = [(i, j) for i in xrange(len(A)) for j in xrange(len(A[0])) if A[i,j][0] == x[0] and A[i,j][1] == x[1]]


This works, but I feel like there should be something more efficient (the arrays I'm working with are much larger).

I tried
np.where()
but that doesn't seem to work too well with higher dimensions. It would return indices for all coordinates where the x-coordinate matched, not the x- and y-coordinates.

Any tips would be appreciated.

Answer

You could use a slightly simpler version of your code:

In [115]: A = [[[1.5, 2.0], [1.0, 2.3], [5.4, 2.3]],
     ...:      [[3.2, 4.4], [2.0, 3.1], [0.0, 2.3]],
     ...:      [[1.0, 2.0], [2.3, 3.4], [4.0, 1.1]]]

In [116]: x = [1., 2.]

In [117]: [(i, j) for i, row in enumerate(A) for j, coor in enumerate(row) if coor == x]
Out[117]: [(2, 0)]

But if the arrays are large, you'd better use a vectorized approach:

In [118]: import numpy as np

In [119]: arr = np.array(A)

In [120]: np.argwhere(np.logical_and(arr[:,:,0] == x[0], arr[:,:,1] == x[1]))
Out[120]: array([[2, 0]], dtype=int64)

Edit: An efficient and elegant way of getting the job done would be:

In [158]: np.argwhere(np.all(arr == x, axis=2))
Out[158]: array([[2, 0]], dtype=int64)