Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1""" 

2Set operations for arrays based on sorting. 

3 

4:Contains: 

5 unique, 

6 isin, 

7 ediff1d, 

8 intersect1d, 

9 setxor1d, 

10 in1d, 

11 union1d, 

12 setdiff1d 

13 

14:Notes: 

15 

16For floating point arrays, inaccurate results may appear due to usual round-off 

17and floating point comparison issues. 

18 

19Speed could be gained in some operations by an implementation of 

20sort(), that can provide directly the permutation vectors, avoiding 

21thus calls to argsort(). 

22 

23To do: Optionally return indices analogously to unique for all functions. 

24 

25:Author: Robert Cimrman 

26 

27""" 

28import functools 

29 

30import numpy as np 

31from numpy.core import overrides 

32 

33 

34array_function_dispatch = functools.partial( 

35 overrides.array_function_dispatch, module='numpy') 

36 

37 

38__all__ = [ 

39 'ediff1d', 'intersect1d', 'setxor1d', 'union1d', 'setdiff1d', 'unique', 

40 'in1d', 'isin' 

41 ] 

42 

43 

44def _ediff1d_dispatcher(ary, to_end=None, to_begin=None): 

45 return (ary, to_end, to_begin) 

46 

47 

48@array_function_dispatch(_ediff1d_dispatcher) 

49def ediff1d(ary, to_end=None, to_begin=None): 

50 """ 

51 The differences between consecutive elements of an array. 

52 

53 Parameters 

54 ---------- 

55 ary : array_like 

56 If necessary, will be flattened before the differences are taken. 

57 to_end : array_like, optional 

58 Number(s) to append at the end of the returned differences. 

59 to_begin : array_like, optional 

60 Number(s) to prepend at the beginning of the returned differences. 

61 

62 Returns 

63 ------- 

64 ediff1d : ndarray 

65 The differences. Loosely, this is ``ary.flat[1:] - ary.flat[:-1]``. 

66 

67 See Also 

68 -------- 

69 diff, gradient 

70 

71 Notes 

72 ----- 

73 When applied to masked arrays, this function drops the mask information 

74 if the `to_begin` and/or `to_end` parameters are used. 

75 

76 Examples 

77 -------- 

78 >>> x = np.array([1, 2, 4, 7, 0]) 

79 >>> np.ediff1d(x) 

80 array([ 1, 2, 3, -7]) 

81 

82 >>> np.ediff1d(x, to_begin=-99, to_end=np.array([88, 99])) 

83 array([-99, 1, 2, ..., -7, 88, 99]) 

84 

85 The returned array is always 1D. 

86 

87 >>> y = [[1, 2, 4], [1, 6, 24]] 

88 >>> np.ediff1d(y) 

89 array([ 1, 2, -3, 5, 18]) 

90 

91 """ 

92 # force a 1d array 

93 ary = np.asanyarray(ary).ravel() 

94 

95 # enforce that the dtype of `ary` is used for the output 

96 dtype_req = ary.dtype 

97 

98 # fast track default case 

99 if to_begin is None and to_end is None: 

100 return ary[1:] - ary[:-1] 

101 

102 if to_begin is None: 

103 l_begin = 0 

104 else: 

105 to_begin = np.asanyarray(to_begin) 

106 if not np.can_cast(to_begin, dtype_req, casting="same_kind"): 

107 raise TypeError("dtype of `to_end` must be compatible " 

108 "with input `ary` under the `same_kind` rule.") 

109 

110 to_begin = to_begin.ravel() 

111 l_begin = len(to_begin) 

112 

113 if to_end is None: 

114 l_end = 0 

115 else: 

116 to_end = np.asanyarray(to_end) 

117 if not np.can_cast(to_end, dtype_req, casting="same_kind"): 

118 raise TypeError("dtype of `to_end` must be compatible " 

119 "with input `ary` under the `same_kind` rule.") 

120 

121 to_end = to_end.ravel() 

122 l_end = len(to_end) 

123 

124 # do the calculation in place and copy to_begin and to_end 

125 l_diff = max(len(ary) - 1, 0) 

126 result = np.empty(l_diff + l_begin + l_end, dtype=ary.dtype) 

127 result = ary.__array_wrap__(result) 

128 if l_begin > 0: 

129 result[:l_begin] = to_begin 

130 if l_end > 0: 

131 result[l_begin + l_diff:] = to_end 

132 np.subtract(ary[1:], ary[:-1], result[l_begin:l_begin + l_diff]) 

133 return result 

134 

135 

136def _unpack_tuple(x): 

137 """ Unpacks one-element tuples for use as return values """ 

138 if len(x) == 1: 

139 return x[0] 

140 else: 

141 return x 

142 

143 

144def _unique_dispatcher(ar, return_index=None, return_inverse=None, 

145 return_counts=None, axis=None): 

146 return (ar,) 

147 

148 

149@array_function_dispatch(_unique_dispatcher) 

150def unique(ar, return_index=False, return_inverse=False, 

151 return_counts=False, axis=None): 

152 """ 

153 Find the unique elements of an array. 

154 

155 Returns the sorted unique elements of an array. There are three optional 

156 outputs in addition to the unique elements: 

157 

158 * the indices of the input array that give the unique values 

159 * the indices of the unique array that reconstruct the input array 

160 * the number of times each unique value comes up in the input array 

161 

162 Parameters 

163 ---------- 

164 ar : array_like 

165 Input array. Unless `axis` is specified, this will be flattened if it 

166 is not already 1-D. 

167 return_index : bool, optional 

168 If True, also return the indices of `ar` (along the specified axis, 

169 if provided, or in the flattened array) that result in the unique array. 

170 return_inverse : bool, optional 

171 If True, also return the indices of the unique array (for the specified 

172 axis, if provided) that can be used to reconstruct `ar`. 

173 return_counts : bool, optional 

174 If True, also return the number of times each unique item appears 

175 in `ar`. 

176 

177 .. versionadded:: 1.9.0 

178 

179 axis : int or None, optional 

180 The axis to operate on. If None, `ar` will be flattened. If an integer, 

181 the subarrays indexed by the given axis will be flattened and treated 

182 as the elements of a 1-D array with the dimension of the given axis, 

183 see the notes for more details. Object arrays or structured arrays 

184 that contain objects are not supported if the `axis` kwarg is used. The 

185 default is None. 

186 

187 .. versionadded:: 1.13.0 

188 

189 Returns 

190 ------- 

191 unique : ndarray 

192 The sorted unique values. 

193 unique_indices : ndarray, optional 

194 The indices of the first occurrences of the unique values in the 

195 original array. Only provided if `return_index` is True. 

196 unique_inverse : ndarray, optional 

197 The indices to reconstruct the original array from the 

198 unique array. Only provided if `return_inverse` is True. 

199 unique_counts : ndarray, optional 

200 The number of times each of the unique values comes up in the 

201 original array. Only provided if `return_counts` is True. 

202 

203 .. versionadded:: 1.9.0 

204 

205 See Also 

206 -------- 

207 numpy.lib.arraysetops : Module with a number of other functions for 

208 performing set operations on arrays. 

209 

210 Notes 

211 ----- 

212 When an axis is specified the subarrays indexed by the axis are sorted. 

213 This is done by making the specified axis the first dimension of the array 

214 (move the axis to the first dimension to keep the order of the other axes) 

215 and then flattening the subarrays in C order. The flattened subarrays are 

216 then viewed as a structured type with each element given a label, with the 

217 effect that we end up with a 1-D array of structured types that can be 

218 treated in the same way as any other 1-D array. The result is that the 

219 flattened subarrays are sorted in lexicographic order starting with the 

220 first element. 

221 

222 Examples 

223 -------- 

224 >>> np.unique([1, 1, 2, 2, 3, 3]) 

225 array([1, 2, 3]) 

226 >>> a = np.array([[1, 1], [2, 3]]) 

227 >>> np.unique(a) 

228 array([1, 2, 3]) 

229 

230 Return the unique rows of a 2D array 

231 

232 >>> a = np.array([[1, 0, 0], [1, 0, 0], [2, 3, 4]]) 

233 >>> np.unique(a, axis=0) 

234 array([[1, 0, 0], [2, 3, 4]]) 

235 

236 Return the indices of the original array that give the unique values: 

237 

238 >>> a = np.array(['a', 'b', 'b', 'c', 'a']) 

239 >>> u, indices = np.unique(a, return_index=True) 

240 >>> u 

241 array(['a', 'b', 'c'], dtype='<U1') 

242 >>> indices 

243 array([0, 1, 3]) 

244 >>> a[indices] 

245 array(['a', 'b', 'c'], dtype='<U1') 

246 

247 Reconstruct the input array from the unique values: 

248 

249 >>> a = np.array([1, 2, 6, 4, 2, 3, 2]) 

250 >>> u, indices = np.unique(a, return_inverse=True) 

251 >>> u 

252 array([1, 2, 3, 4, 6]) 

253 >>> indices 

254 array([0, 1, 4, 3, 1, 2, 1]) 

255 >>> u[indices] 

256 array([1, 2, 6, 4, 2, 3, 2]) 

257 

258 """ 

259 ar = np.asanyarray(ar) 

260 if axis is None: 

261 ret = _unique1d(ar, return_index, return_inverse, return_counts) 

262 return _unpack_tuple(ret) 

263 

264 # axis was specified and not None 

265 try: 

266 ar = np.moveaxis(ar, axis, 0) 

267 except np.AxisError: 

268 # this removes the "axis1" or "axis2" prefix from the error message 

269 raise np.AxisError(axis, ar.ndim) 

270 

271 # Must reshape to a contiguous 2D array for this to work... 

272 orig_shape, orig_dtype = ar.shape, ar.dtype 

273 ar = ar.reshape(orig_shape[0], np.prod(orig_shape[1:], dtype=np.intp)) 

274 ar = np.ascontiguousarray(ar) 

275 dtype = [('f{i}'.format(i=i), ar.dtype) for i in range(ar.shape[1])] 

276 

277 # At this point, `ar` has shape `(n, m)`, and `dtype` is a structured 

278 # data type with `m` fields where each field has the data type of `ar`. 

279 # In the following, we create the array `consolidated`, which has 

280 # shape `(n,)` with data type `dtype`. 

281 try: 

282 if ar.shape[1] > 0: 

283 consolidated = ar.view(dtype) 

284 else: 

285 # If ar.shape[1] == 0, then dtype will be `np.dtype([])`, which is 

286 # a data type with itemsize 0, and the call `ar.view(dtype)` will 

287 # fail. Instead, we'll use `np.empty` to explicitly create the 

288 # array with shape `(len(ar),)`. Because `dtype` in this case has 

289 # itemsize 0, the total size of the result is still 0 bytes. 

290 consolidated = np.empty(len(ar), dtype=dtype) 

291 except TypeError: 

292 # There's no good way to do this for object arrays, etc... 

293 msg = 'The axis argument to unique is not supported for dtype {dt}' 

294 raise TypeError(msg.format(dt=ar.dtype)) 

295 

296 def reshape_uniq(uniq): 

297 n = len(uniq) 

298 uniq = uniq.view(orig_dtype) 

299 uniq = uniq.reshape(n, *orig_shape[1:]) 

300 uniq = np.moveaxis(uniq, 0, axis) 

301 return uniq 

302 

303 output = _unique1d(consolidated, return_index, 

304 return_inverse, return_counts) 

305 output = (reshape_uniq(output[0]),) + output[1:] 

306 return _unpack_tuple(output) 

307 

308 

309def _unique1d(ar, return_index=False, return_inverse=False, 

310 return_counts=False): 

311 """ 

312 Find the unique elements of an array, ignoring shape. 

313 """ 

314 ar = np.asanyarray(ar).flatten() 

315 

316 optional_indices = return_index or return_inverse 

317 

318 if optional_indices: 

319 perm = ar.argsort(kind='mergesort' if return_index else 'quicksort') 

320 aux = ar[perm] 

321 else: 

322 ar.sort() 

323 aux = ar 

324 mask = np.empty(aux.shape, dtype=np.bool_) 

325 mask[:1] = True 

326 mask[1:] = aux[1:] != aux[:-1] 

327 

328 ret = (aux[mask],) 

329 if return_index: 

330 ret += (perm[mask],) 

331 if return_inverse: 

332 imask = np.cumsum(mask) - 1 

333 inv_idx = np.empty(mask.shape, dtype=np.intp) 

334 inv_idx[perm] = imask 

335 ret += (inv_idx,) 

336 if return_counts: 

337 idx = np.concatenate(np.nonzero(mask) + ([mask.size],)) 

338 ret += (np.diff(idx),) 

339 return ret 

340 

341 

342def _intersect1d_dispatcher( 

343 ar1, ar2, assume_unique=None, return_indices=None): 

344 return (ar1, ar2) 

345 

346 

347@array_function_dispatch(_intersect1d_dispatcher) 

348def intersect1d(ar1, ar2, assume_unique=False, return_indices=False): 

349 """ 

350 Find the intersection of two arrays. 

351 

352 Return the sorted, unique values that are in both of the input arrays. 

353 

354 Parameters 

355 ---------- 

356 ar1, ar2 : array_like 

357 Input arrays. Will be flattened if not already 1D. 

358 assume_unique : bool 

359 If True, the input arrays are both assumed to be unique, which 

360 can speed up the calculation. Default is False. 

361 return_indices : bool 

362 If True, the indices which correspond to the intersection of the two 

363 arrays are returned. The first instance of a value is used if there are 

364 multiple. Default is False. 

365 

366 .. versionadded:: 1.15.0 

367 

368 Returns 

369 ------- 

370 intersect1d : ndarray 

371 Sorted 1D array of common and unique elements. 

372 comm1 : ndarray 

373 The indices of the first occurrences of the common values in `ar1`. 

374 Only provided if `return_indices` is True. 

375 comm2 : ndarray 

376 The indices of the first occurrences of the common values in `ar2`. 

377 Only provided if `return_indices` is True. 

378 

379 

380 See Also 

381 -------- 

382 numpy.lib.arraysetops : Module with a number of other functions for 

383 performing set operations on arrays. 

384 

385 Examples 

386 -------- 

387 >>> np.intersect1d([1, 3, 4, 3], [3, 1, 2, 1]) 

388 array([1, 3]) 

389 

390 To intersect more than two arrays, use functools.reduce: 

391 

392 >>> from functools import reduce 

393 >>> reduce(np.intersect1d, ([1, 3, 4, 3], [3, 1, 2, 1], [6, 3, 4, 2])) 

394 array([3]) 

395 

396 To return the indices of the values common to the input arrays 

397 along with the intersected values: 

398 

399 >>> x = np.array([1, 1, 2, 3, 4]) 

400 >>> y = np.array([2, 1, 4, 6]) 

401 >>> xy, x_ind, y_ind = np.intersect1d(x, y, return_indices=True) 

402 >>> x_ind, y_ind 

403 (array([0, 2, 4]), array([1, 0, 2])) 

404 >>> xy, x[x_ind], y[y_ind] 

405 (array([1, 2, 4]), array([1, 2, 4]), array([1, 2, 4])) 

406 

407 """ 

408 ar1 = np.asanyarray(ar1) 

409 ar2 = np.asanyarray(ar2) 

410 

411 if not assume_unique: 

412 if return_indices: 

413 ar1, ind1 = unique(ar1, return_index=True) 

414 ar2, ind2 = unique(ar2, return_index=True) 

415 else: 

416 ar1 = unique(ar1) 

417 ar2 = unique(ar2) 

418 else: 

419 ar1 = ar1.ravel() 

420 ar2 = ar2.ravel() 

421 

422 aux = np.concatenate((ar1, ar2)) 

423 if return_indices: 

424 aux_sort_indices = np.argsort(aux, kind='mergesort') 

425 aux = aux[aux_sort_indices] 

426 else: 

427 aux.sort() 

428 

429 mask = aux[1:] == aux[:-1] 

430 int1d = aux[:-1][mask] 

431 

432 if return_indices: 

433 ar1_indices = aux_sort_indices[:-1][mask] 

434 ar2_indices = aux_sort_indices[1:][mask] - ar1.size 

435 if not assume_unique: 

436 ar1_indices = ind1[ar1_indices] 

437 ar2_indices = ind2[ar2_indices] 

438 

439 return int1d, ar1_indices, ar2_indices 

440 else: 

441 return int1d 

442 

443 

444def _setxor1d_dispatcher(ar1, ar2, assume_unique=None): 

445 return (ar1, ar2) 

446 

447 

448@array_function_dispatch(_setxor1d_dispatcher) 

449def setxor1d(ar1, ar2, assume_unique=False): 

450 """ 

451 Find the set exclusive-or of two arrays. 

452 

453 Return the sorted, unique values that are in only one (not both) of the 

454 input arrays. 

455 

456 Parameters 

457 ---------- 

458 ar1, ar2 : array_like 

459 Input arrays. 

460 assume_unique : bool 

461 If True, the input arrays are both assumed to be unique, which 

462 can speed up the calculation. Default is False. 

463 

464 Returns 

465 ------- 

466 setxor1d : ndarray 

467 Sorted 1D array of unique values that are in only one of the input 

468 arrays. 

469 

470 Examples 

471 -------- 

472 >>> a = np.array([1, 2, 3, 2, 4]) 

473 >>> b = np.array([2, 3, 5, 7, 5]) 

474 >>> np.setxor1d(a,b) 

475 array([1, 4, 5, 7]) 

476 

477 """ 

478 if not assume_unique: 

479 ar1 = unique(ar1) 

480 ar2 = unique(ar2) 

481 

482 aux = np.concatenate((ar1, ar2)) 

483 if aux.size == 0: 

484 return aux 

485 

486 aux.sort() 

487 flag = np.concatenate(([True], aux[1:] != aux[:-1], [True])) 

488 return aux[flag[1:] & flag[:-1]] 

489 

490 

491def _in1d_dispatcher(ar1, ar2, assume_unique=None, invert=None): 

492 return (ar1, ar2) 

493 

494 

495@array_function_dispatch(_in1d_dispatcher) 

496def in1d(ar1, ar2, assume_unique=False, invert=False): 

497 """ 

498 Test whether each element of a 1-D array is also present in a second array. 

499 

500 Returns a boolean array the same length as `ar1` that is True 

501 where an element of `ar1` is in `ar2` and False otherwise. 

502 

503 We recommend using :func:`isin` instead of `in1d` for new code. 

504 

505 Parameters 

506 ---------- 

507 ar1 : (M,) array_like 

508 Input array. 

509 ar2 : array_like 

510 The values against which to test each value of `ar1`. 

511 assume_unique : bool, optional 

512 If True, the input arrays are both assumed to be unique, which 

513 can speed up the calculation. Default is False. 

514 invert : bool, optional 

515 If True, the values in the returned array are inverted (that is, 

516 False where an element of `ar1` is in `ar2` and True otherwise). 

517 Default is False. ``np.in1d(a, b, invert=True)`` is equivalent 

518 to (but is faster than) ``np.invert(in1d(a, b))``. 

519 

520 .. versionadded:: 1.8.0 

521 

522 Returns 

523 ------- 

524 in1d : (M,) ndarray, bool 

525 The values `ar1[in1d]` are in `ar2`. 

526 

527 See Also 

528 -------- 

529 isin : Version of this function that preserves the 

530 shape of ar1. 

531 numpy.lib.arraysetops : Module with a number of other functions for 

532 performing set operations on arrays. 

533 

534 Notes 

535 ----- 

536 `in1d` can be considered as an element-wise function version of the 

537 python keyword `in`, for 1-D sequences. ``in1d(a, b)`` is roughly 

538 equivalent to ``np.array([item in b for item in a])``. 

539 However, this idea fails if `ar2` is a set, or similar (non-sequence) 

540 container: As ``ar2`` is converted to an array, in those cases 

541 ``asarray(ar2)`` is an object array rather than the expected array of 

542 contained values. 

543 

544 .. versionadded:: 1.4.0 

545 

546 Examples 

547 -------- 

548 >>> test = np.array([0, 1, 2, 5, 0]) 

549 >>> states = [0, 2] 

550 >>> mask = np.in1d(test, states) 

551 >>> mask 

552 array([ True, False, True, False, True]) 

553 >>> test[mask] 

554 array([0, 2, 0]) 

555 >>> mask = np.in1d(test, states, invert=True) 

556 >>> mask 

557 array([False, True, False, True, False]) 

558 >>> test[mask] 

559 array([1, 5]) 

560 """ 

561 # Ravel both arrays, behavior for the first array could be different 

562 ar1 = np.asarray(ar1).ravel() 

563 ar2 = np.asarray(ar2).ravel() 

564 

565 # Check if one of the arrays may contain arbitrary objects 

566 contains_object = ar1.dtype.hasobject or ar2.dtype.hasobject 

567 

568 # This code is run when 

569 # a) the first condition is true, making the code significantly faster 

570 # b) the second condition is true (i.e. `ar1` or `ar2` may contain 

571 # arbitrary objects), since then sorting is not guaranteed to work 

572 if len(ar2) < 10 * len(ar1) ** 0.145 or contains_object: 

573 if invert: 

574 mask = np.ones(len(ar1), dtype=bool) 

575 for a in ar2: 

576 mask &= (ar1 != a) 

577 else: 

578 mask = np.zeros(len(ar1), dtype=bool) 

579 for a in ar2: 

580 mask |= (ar1 == a) 

581 return mask 

582 

583 # Otherwise use sorting 

584 if not assume_unique: 

585 ar1, rev_idx = np.unique(ar1, return_inverse=True) 

586 ar2 = np.unique(ar2) 

587 

588 ar = np.concatenate((ar1, ar2)) 

589 # We need this to be a stable sort, so always use 'mergesort' 

590 # here. The values from the first array should always come before 

591 # the values from the second array. 

592 order = ar.argsort(kind='mergesort') 

593 sar = ar[order] 

594 if invert: 

595 bool_ar = (sar[1:] != sar[:-1]) 

596 else: 

597 bool_ar = (sar[1:] == sar[:-1]) 

598 flag = np.concatenate((bool_ar, [invert])) 

599 ret = np.empty(ar.shape, dtype=bool) 

600 ret[order] = flag 

601 

602 if assume_unique: 

603 return ret[:len(ar1)] 

604 else: 

605 return ret[rev_idx] 

606 

607 

608def _isin_dispatcher(element, test_elements, assume_unique=None, invert=None): 

609 return (element, test_elements) 

610 

611 

612@array_function_dispatch(_isin_dispatcher) 

613def isin(element, test_elements, assume_unique=False, invert=False): 

614 """ 

615 Calculates `element in test_elements`, broadcasting over `element` only. 

616 Returns a boolean array of the same shape as `element` that is True 

617 where an element of `element` is in `test_elements` and False otherwise. 

618 

619 Parameters 

620 ---------- 

621 element : array_like 

622 Input array. 

623 test_elements : array_like 

624 The values against which to test each value of `element`. 

625 This argument is flattened if it is an array or array_like. 

626 See notes for behavior with non-array-like parameters. 

627 assume_unique : bool, optional 

628 If True, the input arrays are both assumed to be unique, which 

629 can speed up the calculation. Default is False. 

630 invert : bool, optional 

631 If True, the values in the returned array are inverted, as if 

632 calculating `element not in test_elements`. Default is False. 

633 ``np.isin(a, b, invert=True)`` is equivalent to (but faster 

634 than) ``np.invert(np.isin(a, b))``. 

635 

636 Returns 

637 ------- 

638 isin : ndarray, bool 

639 Has the same shape as `element`. The values `element[isin]` 

640 are in `test_elements`. 

641 

642 See Also 

643 -------- 

644 in1d : Flattened version of this function. 

645 numpy.lib.arraysetops : Module with a number of other functions for 

646 performing set operations on arrays. 

647 

648 Notes 

649 ----- 

650 

651 `isin` is an element-wise function version of the python keyword `in`. 

652 ``isin(a, b)`` is roughly equivalent to 

653 ``np.array([item in b for item in a])`` if `a` and `b` are 1-D sequences. 

654 

655 `element` and `test_elements` are converted to arrays if they are not 

656 already. If `test_elements` is a set (or other non-sequence collection) 

657 it will be converted to an object array with one element, rather than an 

658 array of the values contained in `test_elements`. This is a consequence 

659 of the `array` constructor's way of handling non-sequence collections. 

660 Converting the set to a list usually gives the desired behavior. 

661 

662 .. versionadded:: 1.13.0 

663 

664 Examples 

665 -------- 

666 >>> element = 2*np.arange(4).reshape((2, 2)) 

667 >>> element 

668 array([[0, 2], 

669 [4, 6]]) 

670 >>> test_elements = [1, 2, 4, 8] 

671 >>> mask = np.isin(element, test_elements) 

672 >>> mask 

673 array([[False, True], 

674 [ True, False]]) 

675 >>> element[mask] 

676 array([2, 4]) 

677 

678 The indices of the matched values can be obtained with `nonzero`: 

679 

680 >>> np.nonzero(mask) 

681 (array([0, 1]), array([1, 0])) 

682 

683 The test can also be inverted: 

684 

685 >>> mask = np.isin(element, test_elements, invert=True) 

686 >>> mask 

687 array([[ True, False], 

688 [False, True]]) 

689 >>> element[mask] 

690 array([0, 6]) 

691 

692 Because of how `array` handles sets, the following does not 

693 work as expected: 

694 

695 >>> test_set = {1, 2, 4, 8} 

696 >>> np.isin(element, test_set) 

697 array([[False, False], 

698 [False, False]]) 

699 

700 Casting the set to a list gives the expected result: 

701 

702 >>> np.isin(element, list(test_set)) 

703 array([[False, True], 

704 [ True, False]]) 

705 """ 

706 element = np.asarray(element) 

707 return in1d(element, test_elements, assume_unique=assume_unique, 

708 invert=invert).reshape(element.shape) 

709 

710 

711def _union1d_dispatcher(ar1, ar2): 

712 return (ar1, ar2) 

713 

714 

715@array_function_dispatch(_union1d_dispatcher) 

716def union1d(ar1, ar2): 

717 """ 

718 Find the union of two arrays. 

719 

720 Return the unique, sorted array of values that are in either of the two 

721 input arrays. 

722 

723 Parameters 

724 ---------- 

725 ar1, ar2 : array_like 

726 Input arrays. They are flattened if they are not already 1D. 

727 

728 Returns 

729 ------- 

730 union1d : ndarray 

731 Unique, sorted union of the input arrays. 

732 

733 See Also 

734 -------- 

735 numpy.lib.arraysetops : Module with a number of other functions for 

736 performing set operations on arrays. 

737 

738 Examples 

739 -------- 

740 >>> np.union1d([-1, 0, 1], [-2, 0, 2]) 

741 array([-2, -1, 0, 1, 2]) 

742 

743 To find the union of more than two arrays, use functools.reduce: 

744 

745 >>> from functools import reduce 

746 >>> reduce(np.union1d, ([1, 3, 4, 3], [3, 1, 2, 1], [6, 3, 4, 2])) 

747 array([1, 2, 3, 4, 6]) 

748 """ 

749 return unique(np.concatenate((ar1, ar2), axis=None)) 

750 

751 

752def _setdiff1d_dispatcher(ar1, ar2, assume_unique=None): 

753 return (ar1, ar2) 

754 

755 

756@array_function_dispatch(_setdiff1d_dispatcher) 

757def setdiff1d(ar1, ar2, assume_unique=False): 

758 """ 

759 Find the set difference of two arrays. 

760 

761 Return the unique values in `ar1` that are not in `ar2`. 

762 

763 Parameters 

764 ---------- 

765 ar1 : array_like 

766 Input array. 

767 ar2 : array_like 

768 Input comparison array. 

769 assume_unique : bool 

770 If True, the input arrays are both assumed to be unique, which 

771 can speed up the calculation. Default is False. 

772 

773 Returns 

774 ------- 

775 setdiff1d : ndarray 

776 1D array of values in `ar1` that are not in `ar2`. The result 

777 is sorted when `assume_unique=False`, but otherwise only sorted 

778 if the input is sorted. 

779 

780 See Also 

781 -------- 

782 numpy.lib.arraysetops : Module with a number of other functions for 

783 performing set operations on arrays. 

784 

785 Examples 

786 -------- 

787 >>> a = np.array([1, 2, 3, 2, 4, 1]) 

788 >>> b = np.array([3, 4, 5, 6]) 

789 >>> np.setdiff1d(a, b) 

790 array([1, 2]) 

791 

792 """ 

793 if assume_unique: 

794 ar1 = np.asarray(ar1).ravel() 

795 else: 

796 ar1 = unique(ar1) 

797 ar2 = unique(ar2) 

798 return ar1[in1d(ar1, ar2, assume_unique=True, invert=True)]