diff options
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/arraysetops.py | 62 | ||||
| -rw-r--r-- | numpy/lib/tests/test_arraysetops.py | 46 |
2 files changed, 89 insertions, 19 deletions
diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py index 9afb00f29..c5e7822f2 100644 --- a/numpy/lib/arraysetops.py +++ b/numpy/lib/arraysetops.py @@ -8,6 +8,7 @@ Set operations for 1D numeric arrays based on sorting. intersect1d_nu, setxor1d, setmember1d, + setmember1d_nu, union1d, setdiff1d @@ -33,7 +34,7 @@ last revision: 07.01.2007 :Author: Robert Cimrman """ __all__ = ['ediff1d', 'unique1d', 'intersect1d', 'intersect1d_nu', 'setxor1d', - 'setmember1d', 'union1d', 'setdiff1d'] + 'setmember1d', 'setmember1d_nu', 'union1d', 'setdiff1d'] import numpy as np @@ -287,6 +288,7 @@ def setmember1d(ar1, ar2): See Also -------- + setmember1d_nu : Works for arrays with non-unique elements. numpy.lib.arraysetops : Module with a number of other functions for performing set operations on arrays. @@ -301,30 +303,52 @@ def setmember1d(ar1, ar2): array([0, 2]) """ - ar1 = np.asarray( ar1 ) - ar2 = np.asarray( ar2 ) - ar = np.concatenate( (ar1, ar2 ) ) - b1 = np.zeros( ar1.shape, dtype = np.int8 ) - b2 = np.ones( ar2.shape, dtype = np.int8 ) - tt = np.concatenate( (b1, b2) ) - # We need this to be a stable sort, so always use 'mergesort' here. The # values from the first array should always come before the values from the # second array. - perm = ar.argsort(kind='mergesort') - aux = ar[perm] - aux2 = tt[perm] -# flag = ediff1d( aux, 1 ) == 0 - flag = np.concatenate( (aux[1:] == aux[:-1], [False] ) ) - ii = np.where( flag * aux2 )[0] - aux = perm[ii+1] - perm[ii+1] = perm[ii] - perm[ii] = aux - - indx = perm.argsort(kind='mergesort')[:len( ar1 )] + ar = np.concatenate( (ar1, ar2 ) ) + order = ar.argsort(kind='mergesort') + sar = ar[order] + equal_adj = (sar[1:] == sar[:-1]) + flag = np.concatenate( (equal_adj, [False] ) ) + indx = order.argsort(kind='mergesort')[:len( ar1 )] return flag[indx] +def setmember1d_nu(ar1, ar2): + """ + Return a boolean array set True where first element is in second array. + + Boolean array is the shape of `ar1` containing True where the elements + of `ar1` are in `ar2` and False otherwise. + + Unlike setmember1d(), this version works also for arrays with duplicate + values. It uses setmember1d() internally. For arrays with unique + entries it is slower than calling setmember1d() directly. + + Parameters + ---------- + ar1 : array_like + Input array. + ar2 : array_like + Input array. + + Returns + ------- + mask : ndarray, bool + The values `ar1[mask]` are in `ar2`. + + See Also + -------- + setmember1d : Faster for arrays with unique elements. + numpy.lib.arraysetops : Module with a number of other functions for + performing set operations on arrays. + + """ + unique_ar1, rev_idx = np.unique1d(ar1, return_inverse=True) + mask = np.setmember1d(unique_ar1, np.unique1d(ar2)) + return mask[rev_idx] + def union1d(ar1, ar2): """ Union of 1D arrays with unique elements. diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py index c40cf9d20..40bc11f6e 100644 --- a/numpy/lib/tests/test_arraysetops.py +++ b/numpy/lib/tests/test_arraysetops.py @@ -102,6 +102,52 @@ class TestAso(TestCase): assert_array_equal([], setmember1d([],[])) + def test_setmember1d_nu(self): + a = np.array([5,4,5,3,4,4,3,4,3,5,2,1,5,5]) + b = [2,3,4] + + ec = [False, True, False, True, True, True, True, True, True, False, + True, False, False, False] + c = setmember1d_nu(a, b) + assert_array_equal(c, ec) + + b = b + [5, 5, 4] + + ec = [True, True, True, True, True, True, True, True, True, True, + True, False, True, True] + c = setmember1d_nu(a, b) + assert_array_equal(c, ec) + + a = np.array([5, 7, 1, 2]) + b = np.array([2, 4, 3, 1, 5]) + + ec = np.array([True, False, True, True]) + c = setmember1d_nu(a, b) + assert_array_equal(c, ec) + + a = np.array([5, 7, 1, 1, 2]) + b = np.array([2, 4, 3, 3, 1, 5]) + + ec = np.array([True, False, True, True, True]) + c = setmember1d_nu(a, b) + assert_array_equal(c, ec) + + a = np.array([5]) + b = np.array([2]) + + ec = np.array([False]) + c = setmember1d_nu(a, b) + assert_array_equal(c, ec) + + a = np.array([5, 5]) + b = np.array([2, 2]) + + ec = np.array([False, False]) + c = setmember1d_nu(a, b) + assert_array_equal(c, ec) + + assert_array_equal(setmember1d_nu([], []), []) + def test_union1d( self ): a = np.array( [5, 4, 7, 1, 2] ) b = np.array( [2, 4, 3, 3, 2, 1, 5] ) |