diff options
| -rw-r--r-- | numpy/lib/arraysetops.py | 20 | ||||
| -rw-r--r-- | numpy/lib/tests/test_arraysetops.py | 31 |
2 files changed, 41 insertions, 10 deletions
diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py index fa332e7fd..cf5f47a82 100644 --- a/numpy/lib/arraysetops.py +++ b/numpy/lib/arraysetops.py @@ -621,28 +621,28 @@ def in1d(ar1, ar2, assume_unique=False, invert=False, *, kind=None): # Ensure that iteration through object arrays yields size-1 arrays if ar2.dtype == object: ar2 = ar2.reshape(-1, 1) - # Convert booleans to uint8 so we can use the fast integer algorithm - if ar1.dtype == bool: - ar1 = ar1 + np.uint8(0) - if ar2.dtype == bool: - ar2 = ar2 + np.uint8(0) - - # Check if we can use a fast integer algorithm: - integer_arrays = (np.issubdtype(ar1.dtype, np.integer) and - np.issubdtype(ar2.dtype, np.integer)) if kind not in {None, 'sort', 'table'}: raise ValueError( f"Invalid kind: '{kind}'. Please use None, 'sort' or 'table'.") - if integer_arrays and kind in {None, 'table'}: + # Can use the table method if all arrays are integers or boolean: + is_int_arrays = all(ar.dtype.kind in ("u", "i", "b") for ar in (ar1, ar2)) + use_table_method = is_int_arrays and kind in {None, 'table'} + if use_table_method: if ar2.size == 0: if invert: return np.ones_like(ar1, dtype=bool) else: return np.zeros_like(ar1, dtype=bool) + # Convert booleans to uint8 so we can use the fast integer algorithm + if ar1.dtype == bool: + ar1 = ar1.astype(np.uint8) + if ar2.dtype == bool: + ar2 = ar2.astype(np.uint8) + ar2_min = np.min(ar2) ar2_max = np.max(ar2) diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py index a21e5f640..bb07e25a9 100644 --- a/numpy/lib/tests/test_arraysetops.py +++ b/numpy/lib/tests/test_arraysetops.py @@ -396,6 +396,37 @@ class TestSetOps: assert_array_equal(np.invert(expected), in1d(a, b, invert=True, kind=kind)) + @pytest.mark.parametrize("kind", [None, "sort"]) + def test_in1d_timedelta(self, kind): + """Test that in1d works for timedelta input""" + rstate = np.random.RandomState(0) + a = rstate.randint(0, 100, size=10) + b = rstate.randint(0, 100, size=10) + truth = in1d(a, b) + a_timedelta = a.astype("timedelta64[s]") + b_timedelta = b.astype("timedelta64[s]") + assert_array_equal(truth, in1d(a_timedelta, b_timedelta, kind=kind)) + + def test_in1d_table_timedelta_fails(self): + a = np.array([0, 1, 2], dtype="timedelta64[s]") + b = a + # Make sure it raises a value error: + with pytest.raises(ValueError): + in1d(a, b, kind="table") + + @pytest.mark.parametrize("kind", [None, "sort", "table"]) + def test_in1d_mixed_boolean(self, kind): + """Test that in1d works as expected for bool/int input.""" + for dtype in np.typecodes["AllInteger"]: + a = np.array([True, False, False], dtype=bool) + b = np.array([1, 1, 1, 1], dtype=dtype) + expected = np.array([True, False, False], dtype=bool) + assert_array_equal(in1d(a, b, kind=kind), expected) + + a, b = b, a + expected = np.array([True, True, True, True], dtype=bool) + assert_array_equal(in1d(a, b, kind=kind), expected) + def test_in1d_first_array_is_object(self): ar1 = [None] ar2 = np.array([1]*10) |