Take the 2-minute tour ×
Stack Overflow is a question and answer site for professional and enthusiast programmers. It's 100% free, no registration required.

Here is a small code to illustrate the problem.

A = array([[1,2], [1,0], [5,3]])
f_of_A = f(A)   # this is precomputed and expensive


values = array([[1,2], [1,0]])


# location of values in A
# if I just had 1d values I could use numpy.in1d here
indices = array([0, 1])


# example of operation type I need (recalculating f_of_A as needed is not an option)
f_of_A[ indices ]

So, basically I think I need some equivalent to in1d for higher dimensions. Does such a thing exist? Or is there some other approach?

Looks like there is also a searchsorted() function, but that seems to work for 1d arrays also. In this example I used 2d points, but any solution would need to work for 3d points also.

share|improve this question

2 Answers 2

up vote 2 down vote accepted

Okay, this is what I came up with.

To find the value of one multi-dimensional index, let's say ii = np.array([1,2]), we can do:

n.where((A == ii).all(axis=1))[0]

Let's break this down, we have A == ii, which will give element-wise comparisons with ii for each row of A. We want an entire row to be true, so we add .all(axis=1) to collapse them. To find where these indices happen, we plug this into np.where and get the first value of the tuple.

Now, I don't have a fast way to do this with multiple indices yet (although I have a feeling there is one). However, this will get the job done:

np.hstack([np.where((A == values[i]).all(axis=1))[0] for i in xrange(len(values))])

This basically just calls the above, for each value of values, and concatenates the result.

Update:

Here is for the multi-dimensional case (all in one go, should be fairly fast):

np.where((np.expand_dims(A, -1) == values.T).all(axis=1).any(axis=1))[0]
share|improve this answer
    
Could you explain how axis works? I was reading about np.all and didn't understand what they meant by axis –  A.Wan Jul 22 '13 at 20:58

You can use np.in1d over a view of your original array with all coordinates collapsed into a single variable of dtype np.void:

import numpy as np

A = np.array([[1,2], [1,0], [5,3]])
values = np.array([[1,2], [1,0]])

# Make sure both arrays are contiguous and have common dtype
common_dtype = np.common_type(A, values)
a = np.ascontiguousarray(A, dtype=common_dtype)
vals = np.ascontiguousarray(values, dtype=common_dtype)

a_view = A.view((np.void, A.dtype.itemsize*A.shape[1])).ravel()
values_view = values.view((np.void,
                           values.dtype.itemsize*values.shape[1])).ravel()

Now each item of a_view and values_view is all coordinates for one point packed together, so you can do whatever 1D magic you would use. I don't see how to use np.in1d to find indices though, so I would go the np.searchsorted route:

sort_idx = np.argsort(a_view)
locations = np.searchsorted(a_view, values_view, sorter=sort_idx)
locations = sort_idx[locations]

>>> locations
array([0, 1], dtype=int64)
share|improve this answer
    
Thanks, I attempted using a 1d view, but I think my problem was not using np.ascontiguousarray(). As for the np.in1d this would be to generate a mask instead of the indices –  user1984528 Jul 22 '13 at 21:10

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

Not the answer you're looking for? Browse other questions tagged or ask your own question.