CODEWITHSUNDEEP
pythonarrayspython-2.7numpy
berkelem
berkelem

Reputation: 2125

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.

Upvotes: 0

Views: 3480

Answers (3)

Tonechas
Tonechas

Reputation: 13743

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)

Upvotes: 2

alexis
alexis

Reputation: 50220

If you'll be doing a lot of lookups, you can build an index mapping the coordinates to positions in your matrix. E.g., (1.5, 2.0) maps to (0,0), (1.0, 2.3) maps to (0, 1), etc. Here's how to build an index of your array:

>>> 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]]]
>>> revind = dict()
>>> for r, row in enumerate(A):
...     for c, pt in enumerate(row):
...         revind[tuple(pt)] = (r, c)
... 
>>> revind[(1.0, 2.0)]
(2,0)

Or as a (less readable) comprehension:

>>> revind = dict((tuple(pt), (r,c)) for r,row in enumerate(A) for c,pt in enumerate(row))

Note that you have to use tuples, since lists are "mutable" and cannot be used as dictionary keys. Doing it this way ensures that each lookup takes place in O(1) time, which is dramatically faster than searching a long list (as in your question and in the accepted answer).

Upvotes: 1

akuiper
akuiper

Reputation: 215117

You can possibly try this:

import numpy as np
arrA = np.array(A)
x = [1.0, 2.0]

np.where((arrA == x).sum(axis = 2) == 2)
# (array([2]), array([0]))

Upvotes: 1

Related Questions