diff options
| -rw-r--r-- | numpy/lib/arraysetops.py | 12 | ||||
| -rw-r--r-- | numpy/lib/tests/test_arraysetops.py | 102 |
2 files changed, 59 insertions, 55 deletions
diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py index 721039238..98597bc46 100644 --- a/numpy/lib/arraysetops.py +++ b/numpy/lib/arraysetops.py @@ -277,7 +277,7 @@ def setxor1d(ar1, ar2, assume_unique=False): def in1d(ar1, ar2, assume_unique=False): """ - Test whether each element of a 1D array is also present in a second array. + Test whether each element of a 1-D array is also present in a second array. Returns a boolean array the same length as `ar1` that is True where an element of `ar1` is in `ar2` and False otherwise. @@ -305,7 +305,7 @@ def in1d(ar1, ar2, assume_unique=False): Notes ----- `in1d` can be considered as an element-wise function version of the - python keyword `in`, for 1D sequences. ``in1d(a, b)`` is roughly + python keyword `in`, for 1-D sequences. ``in1d(a, b)`` is roughly equivalent to ``np.array([item in b for item in a])``. .. versionadded:: 1.4.0 @@ -321,6 +321,14 @@ def in1d(ar1, ar2, assume_unique=False): array([0, 2, 0]) """ + # 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) + return mask + + # Otherwise use sorting if not assume_unique: ar1, rev_idx = np.unique(ar1, return_inverse=True) ar2 = np.unique(ar2) diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py index 907a27a8c..e40c155a4 100644 --- a/numpy/lib/tests/test_arraysetops.py +++ b/numpy/lib/tests/test_arraysetops.py @@ -90,66 +90,62 @@ class TestAso(TestCase): assert_array_equal([1],ediff1d(two_elem)) def test_in1d(self): - a = np.array( [5, 7, 1, 2] ) - b = np.array( [2, 4, 3, 1, 5] ) - - ec = np.array( [True, False, True, True] ) - c = in1d( a, b, assume_unique=True ) - assert_array_equal( c, ec ) - - a[0] = 8 - ec = np.array( [False, False, True, True] ) - c = in1d( a, b, assume_unique=True ) - assert_array_equal( c, ec ) - - a[0], a[3] = 4, 8 - ec = np.array( [True, False, True, False] ) - c = in1d( a, b, assume_unique=True ) - assert_array_equal( c, ec ) - - 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 = in1d(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 = in1d(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 = in1d(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 = in1d(a, b) - assert_array_equal(c, ec) + # 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, 7, 1, 2]) + b = np.array([2, 4, 3, 1, 5] * mult) + ec = np.array([True, False, True, True]) + c = in1d(a, b, assume_unique=True) + assert_array_equal(c, ec) + + a[0] = 8 + ec = np.array([False, False, True, True]) + c = in1d(a, b, assume_unique=True) + assert_array_equal(c, ec) + + a[0], a[3] = 4, 8 + ec = np.array([True, False, True, False]) + c = in1d(a, b, assume_unique=True) + assert_array_equal(c, ec) + + a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5]) + b = [2, 3, 4] * mult + ec = [False, True, False, True, True, True, True, True, True, False, + True, False, False, False] + c = in1d(a, b) + assert_array_equal(c, ec) + + b = b + [5, 5, 4] * mult + ec = [True, True, True, True, True, True, True, True, True, True, + True, False, True, True] + c = in1d(a, b) + assert_array_equal(c, ec) + + a = np.array([5, 7, 1, 2]) + b = np.array([2, 4, 3, 1, 5] * mult) + ec = np.array([True, False, True, True]) + c = in1d(a, b) + assert_array_equal(c, ec) + + a = np.array([5, 7, 1, 1, 2]) + b = np.array([2, 4, 3, 3, 1, 5] * mult) + ec = np.array([True, False, True, True, True]) + c = in1d(a, b) + assert_array_equal(c, ec) + + a = np.array([5, 5]) + b = np.array([2, 2] * mult) + ec = np.array([False, False]) + c = in1d(a, b) + assert_array_equal(c, ec) a = np.array([5]) b = np.array([2]) - ec = np.array([False]) c = in1d(a, b) assert_array_equal(c, ec) - a = np.array([5, 5]) - b = np.array([2, 2]) - - ec = np.array([False, False]) - c = in1d(a, b) - assert_array_equal(c, ec) - assert_array_equal(in1d([], []), []) def test_in1d_char_array( self ): |