diff options
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/arraysetops.py | 33 | ||||
| -rw-r--r-- | numpy/lib/tests/test_arraysetops.py | 9 |
2 files changed, 35 insertions, 7 deletions
diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py index a8b0a95bb..5cd535703 100644 --- a/numpy/lib/arraysetops.py +++ b/numpy/lib/arraysetops.py @@ -281,7 +281,7 @@ def setxor1d(ar1, ar2, assume_unique=False): flag2 = flag[1:] == flag[:-1] return aux[flag2] -def in1d(ar1, ar2, assume_unique=False): +def in1d(ar1, ar2, assume_unique=False, invert=False): """ Test whether each element of a 1-D array is also present in a second array. @@ -297,6 +297,13 @@ def in1d(ar1, ar2, assume_unique=False): assume_unique : bool, optional If True, the input arrays are both assumed to be unique, which can speed up the calculation. Default is False. + invert : bool, optional + If True, the values in the returned array are inverted (that is, + False where an element of `ar1` is in `ar2` and True otherwise). + Default is False. ``np.in1d(a, b, invert=True)`` is equivalent + to (but is faster than) ``np.invert(in1d(a, b))``. + + .. versionadded:: 1.8.0 Returns ------- @@ -325,7 +332,11 @@ def in1d(ar1, ar2, assume_unique=False): array([ True, False, True, False, True], dtype=bool) >>> test[mask] array([0, 2, 0]) - + >>> mask = np.in1d(test, states, invert=True) + >>> mask + array([False, True, False, True, False], dtype=bool) + >>> test[mask] + array([1, 5]) """ # Ravel both arrays, behavior for the first array could be different ar1 = np.asarray(ar1).ravel() @@ -333,9 +344,14 @@ def in1d(ar1, ar2, assume_unique=False): # This code is significantly faster when the condition is satisfied. if len(ar2) < 10 * len(ar1) ** 0.145: - mask = np.zeros(len(ar1), dtype=np.bool) - for a in ar2: - mask |= (ar1 == a) + if invert: + mask = np.ones(len(ar1), dtype=np.bool) + for a in ar2: + mask &= (ar1 != a) + else: + mask = np.zeros(len(ar1), dtype=np.bool) + for a in ar2: + mask |= (ar1 == a) return mask # Otherwise use sorting @@ -349,8 +365,11 @@ def in1d(ar1, ar2, assume_unique=False): # the values from the second array. order = ar.argsort(kind='mergesort') sar = ar[order] - equal_adj = (sar[1:] == sar[:-1]) - flag = np.concatenate( (equal_adj, [False] ) ) + if invert: + bool_ar = (sar[1:] != sar[:-1]) + else: + bool_ar = (sar[1:] == sar[:-1]) + flag = np.concatenate( (bool_ar, [invert]) ) indx = order.argsort(kind='mergesort')[:len( ar1 )] if assume_unique: diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py index b3c608cf8..f7f75922d 100644 --- a/numpy/lib/tests/test_arraysetops.py +++ b/numpy/lib/tests/test_arraysetops.py @@ -190,6 +190,15 @@ class TestSetOps(TestCase): assert_array_equal(c, ec) + def test_in1d_invert(self): + "Test in1d's invert parameter" + # We use two different sizes for the b array here to test the + # two different paths in in1d(). + for mult in (1, 10): + a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5]) + b = [2, 3, 4] * mult + assert_array_equal(np.invert(in1d(a, b)), in1d(a, b, invert=True)) + def test_in1d_ravel(self): # Test that in1d ravels its input arrays. This is not documented # behavior however. The test is to ensure consistentency. |