test for membership in a 2d numpy array
What we'd really like to do is use np.in1d
... except that np.in1d
only works with 1-dimensional arrays. Our arrays are multi-dimensional. However, we can view the arrays as a 1-dimensional array of strings:
arr.view(np.dtype((np.void, arr.dtype.itemsize * arr.shape[-1])))
For example,
In [15]: arr = np.array([[1, 2], [2, 3], [1, 3]])
In [16]: arr = arr.view(np.dtype((np.void, arr.dtype.itemsize * arr.shape[-1])))
In [30]: arr.dtype
Out[30]: dtype('V16')
In [31]: arr.shape
Out[31]: (3, 1)
In [37]: arr
Out[37]:
array([[b'\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00'],
[b'\x02\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00'],
[b'\x01\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00']],
dtype='|V16')
This makes each row of arr
a string. Now it is just a matter of hooking this up
to np.in1d
:
import numpy as np
def asvoid(arr):
"""
Based on http://stackoverflow.com/a/16973510/190597 (Jaime, 2013-06)
View the array as dtype np.void (bytes). The items along the last axis are
viewed as one value. This allows comparisons to be performed on the entire row.
"""
arr = np.ascontiguousarray(arr)
if np.issubdtype(arr.dtype, np.floating):
""" Care needs to be taken here since
np.array([-0.]).view(np.void) != np.array([0.]).view(np.void)
Adding 0. converts -0. to 0.
"""
arr += 0.
return arr.view(np.dtype((np.void, arr.dtype.itemsize * arr.shape[-1])))
def inNd(a, b, assume_unique=False):
a = asvoid(a)
b = asvoid(b)
return np.in1d(a, b, assume_unique)
tests = [
(np.array([[1, 2], [2, 3], [1, 3]]),
np.array([[2, 2], [3, 3], [4, 4]]),
np.array([False, False, False])),
(np.array([[1, 2], [2, 2], [1, 3]]),
np.array([[2, 2], [3, 3], [4, 4]]),
np.array([True, False, False])),
(np.array([[1, 2], [3, 4], [5, 6]]),
np.array([[1, 2], [3, 4], [7, 8]]),
np.array([True, True, False])),
(np.array([[1, 2], [5, 6], [3, 4]]),
np.array([[1, 2], [5, 6], [7, 8]]),
np.array([True, True, False])),
(np.array([[-0.5, 2.5, -2, 100, 2], [5, 6, 7, 8, 9], [3, 4, 5, 6, 7]]),
np.array([[1.0, 2, 3, 4, 5], [5, 6, 7, 8, 9], [-0.5, 2.5, -2, 100, 2]]),
np.array([False, True, True]))
]
for a, b, answer in tests:
result = inNd(b, a)
try:
assert np.all(answer == result)
except AssertionError:
print('''\
a:
{a}
b:
{b}
answer: {answer}
result: {result}'''.format(**locals()))
raise
else:
print('Success!')
yields
Success!
In [1]: import numpy as np
In [2]: a = np.array([[1,2],[3,4]])
In [3]: b = np.array([[3,4],[1,2]])
In [5]: a = a[a[:,1].argsort(kind='mergesort')]
In [6]: a = a[a[:,0].argsort(kind='mergesort')]
In [7]: b = b[b[:,1].argsort(kind='mergesort')]
In [8]: b = b[b[:,0].argsort(kind='mergesort')]
In [9]: bInA1 = b[:,0] == a[:,0]
In [10]: bInA2 = b[:,1] == a[:,1]
In [11]: bInA = bInA1*bInA2
In [12]: bInA
Out[12]: array([ True, True], dtype=bool)
should do this... Not sure, whether this is still efficient. You need do mergesort
, as other methods are unstable.
Edit:
If you have more than 2 columns and if the rows are sorted already, you can do
In [24]: bInA = np.array([True,]*a.shape[0])
In [25]: bInA
Out[25]: array([ True, True], dtype=bool)
In [26]: for k in range(a.shape[1]):
bInAk = b[:,k] == a[:,k]
bInA = bInAk*bInA
....:
In [27]: bInA
Out[27]: array([ True, True], dtype=bool)
There is still space for speeding up, as in the iteration, you don't have to check the entire column, but only the entries where the current bInA
is True
.
If you have smth like a=np.array([[1,2],[3,4],[5,6]])
and b=np.array([[5,6],[1,2],[7,6]])
, you can convert them into complex 1-D array:
c=a[:,0]+a[:,1]*1j
d=b[:,0]+b[:,1]*1j
This whole stuff in my Interpreter looks like this:
>>> c=a[:,0]+a[:,1]*1j
>>> c
array([ 1.+2.j, 3.+4.j, 5.+6.j])
>>> d=b[:,0]+b[:,1]*1j
>>> d
array([ 5.+6.j, 1.+2.j, 7.+6.j])
And now that you have just 1D array, you can easily do np.in1d(c,d)
, and the Python will give you:
>>> np.in1d(c,d)
array([ True, False, True], dtype=bool)
And with this you don't need any loops, at least with this data type