ENH: speed up in1d() in the case of ar1 >> ar2. Closes #1603.

A timing script justifying the switching criterion is attached to ticket 1603.

Thanks to Neil Crighton.

1 files changed, 49 insertions, 53 deletions

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 ):