@@ -571,28 +571,28 @@ def unique_values(x):
571
571
572
572
573
573
def _intersect1d_dispatcher (
574
- ar1 , ar2 , assume_unique = None , return_indices = None ):
575
- return ( ar1 , ar2 )
574
+ * ars , assume_unique = None , return_indices = None ):
575
+ return ars
576
576
577
577
578
578
@array_function_dispatch (_intersect1d_dispatcher )
579
- def intersect1d (ar1 , ar2 , assume_unique = False , return_indices = False ):
579
+ def intersect1d (* ars , assume_unique = False , return_indices = False ):
580
580
"""
581
- Find the intersection of two arrays.
581
+ Find the intersection of multiple arrays.
582
582
583
- Return the sorted, unique values that are in both of the input arrays.
583
+ Return the sorted, unique values that are in all of the input arrays.
584
584
585
585
Parameters
586
586
----------
587
- ar1, ar2 : array_like
588
- Input arrays. Will be flattened if not already 1D.
587
+ *ars : array_like
588
+ Input arrays. Each will be flattened if not already 1D.
589
589
assume_unique : bool
590
- If True, the input arrays are both assumed to be unique, which
591
- can speed up the calculation. If True but ``ar1`` or ``ar2`` are not
592
- unique, incorrect results and out-of-bounds indices could result.
590
+ If True, the input arrays are all assumed to be unique, which
591
+ can speed up the calculation. If True but any of the arrays in ars is
592
+ not unique, incorrect results and out-of-bounds indices could result.
593
593
Default is False.
594
594
return_indices : bool
595
- If True, the indices which correspond to the intersection of the two
595
+ If True, the indices which correspond to the intersection of all
596
596
arrays are returned. The first instance of a value is used if there are
597
597
multiple. Default is False.
598
598
@@ -602,23 +602,28 @@ def intersect1d(ar1, ar2, assume_unique=False, return_indices=False):
602
602
-------
603
603
intersect1d : ndarray
604
604
Sorted 1D array of common and unique elements.
605
- comm1 : ndarray
606
- The indices of the first occurrences of the common values in `ar1`.
607
- Only provided if `return_indices` is True.
608
- comm2 : ndarray
609
- The indices of the first occurrences of the common values in `ar2`.
605
+ *comms : list of ndarray
606
+ The indices of the first occurrences of the common values in `ars`.
610
607
Only provided if `return_indices` is True.
608
+ comms[0] contains the indices for ars[0],
609
+ comms[1] contains the indices for ars[1] and so on
610
+
611
+
612
+ See Also
613
+ --------
614
+ numpy.lib.arraysetops : Module with a number of other functions for
615
+ performing set operations on arrays.
611
616
612
617
Examples
613
618
10000
--------
614
619
>>> np.intersect1d([1, 3, 4, 3], [3, 1, 2, 1])
615
620
array([1, 3])
616
621
617
- To intersect more than two arrays, use functools.reduce :
622
+ To intersect more than two arrays, use:
618
623
619
- >>> from functools import reduce
620
- >>> reduce(np. intersect1d, ([1, 3, 4, 3], [3, 1, 2, 1], [6, 3, 4, 2]) )
621
- array([3])
624
+ >>> ars = ([1, 3, 4, 3], [3, 1, 2, 1], [6, 3, 4, 2])
625
+ >>> intersect1d(*ars, return_indices=True )
626
+ ( array([3]), array([1]), array([0]), array([1]) )
622
627
623
628
To return the indices of the values common to the input arrays
624
629
along with the intersected values:
@@ -632,38 +637,43 @@ def intersect1d(ar1, ar2, assume_unique=False, return_indices=False):
632
637
(array([1, 2, 4]), array([1, 2, 4]), array([1, 2, 4]))
633
638
634
639
"""
635
- ar1 = np .asanyarray (ar1 )
636
- ar2 = np .asanyarray (ar2 )
640
+ ars = [np .asanyarray (ar ) for ar in ars ]
637
641
638
642
if not assume_unique :
639
643
if return_indices :
640
- ar1 , ind1 = unique (ar1 , return_index = True )
641
- ar2 , ind2 = unique (ar2 , return_index = True )
644
+ inds = [None ] * len (ars )
645
+ for i , ar in enumerate (ars ):
646
+ ars [i ], inds [i ] = unique (ar , return_index = True )
642
647
else :
643
- ar1 = unique ( ar1 )
644
- ar2 = unique (ar2 )
648
+ for i , ar in enumerate ( ars ):
649
+ ars [ i ] = unique (ar )
645
650
else :
646
- ar1 = ar1 . ravel ()
647
- ar2 = ar2 .ravel ()
651
+ for i , ar in enumerate ( ars ):
652
+ ars [ i ] = ar .ravel ()
648
653
649
- aux = np .concatenate (( ar1 , ar2 ) )
654
+ aux = np .concatenate (ars )
650
655
if return_indices :
651
656
aux_sort_indices = np .argsort (aux , kind = 'mergesort' )
652
657
aux = aux [aux_sort_indices ]
653
658
else :
654
659
aux .sort ()
655
660
656
- mask = aux [1 :] == aux [:- 1 ]
657
- int1d = aux [:- 1 ][mask ]
661
+ # aux is sorted and each array in ars has only unique elements.
662
+ # The same element in a distance of len(ars)+1 away means,
663
+ # that the element must have been in each of the arrays in ars.
664
+ mask = aux [:- len (ars )+ 1 ] == aux [len (ars )- 1 :]
665
+ int1d = aux [:- len (ars )+ 1 ][mask ]
658
666
659
667
if return_indices :
660
- ar1_indices = aux_sort_indices [:- 1 ][mask ]
661
- ar2_indices = aux_sort_indices [1 :][mask ] - ar1 .size
662
- if not assume_unique :
663
- ar1_indices = ind1 [ar1_indices ]
664
- ar2_indices = ind2 [ar2_indices ]
665
-
666
- return int1d , ar1_indices , ar2_indices
668
+ ret_indizes = [None ] * len (ars )
669
+ offset = 0
670
+ for i , ar in enumerate (ars ):
671
+ imax = aux_sort_indices .size - len (ars ) + i + 1
672
+ ret_indizes [i ] = aux_sort_indices [i :imax ][mask ] - offset
673
+ offset += ar .size
674
+ if not assume_unique :
675
+ ret_indizes [i ] = inds [i ][ret_indizes [i ]]
676
+ return int1d , * ret_indizes
667
677
else :
668
678
return int1d
669
679
0 commit comments