diff options
Diffstat (limited to 'numpy/lib/arraysetops.py')
| -rw-r--r-- | numpy/lib/arraysetops.py | 26 |
1 files changed, 21 insertions, 5 deletions
diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py index e63da519b..7776d7e76 100644 --- a/numpy/lib/arraysetops.py +++ b/numpy/lib/arraysetops.py @@ -205,8 +205,9 @@ def unique(ar, return_index=False, return_inverse=False, return_counts=False): ret += (perm[flag],) if return_inverse: iflag = np.cumsum(flag) - 1 - iperm = perm.argsort() - ret += (np.take(iflag, iperm),) + inv_idx = np.empty(ar.shape, dtype=np.intp) + inv_idx[perm] = iflag + ret += (inv_idx,) if return_counts: idx = np.concatenate(np.nonzero(flag) + ([ar.size],)) ret += (np.diff(idx),) @@ -241,6 +242,11 @@ def intersect1d(ar1, ar2, assume_unique=False): >>> np.intersect1d([1, 3, 4, 3], [3, 1, 2, 1]) array([1, 3]) + To intersect more than two arrays, use functools.reduce: + + >>> from functools import reduce + >>> reduce(np.intersect1d, ([1, 3, 4, 3], [3, 1, 2, 1], [6, 3, 4, 2])) + array([3]) """ if not assume_unique: # Might be faster than unique( intersect1d( ar1, ar2 ) )? @@ -333,6 +339,10 @@ def in1d(ar1, ar2, assume_unique=False, invert=False): `in1d` can be considered as an element-wise function version of the python keyword `in`, for 1-D sequences. ``in1d(a, b)`` is roughly equivalent to ``np.array([item in b for item in a])``. + However, this idea fails if `ar2` is a set, or similar (non-sequence) + container: As ``ar2`` is converted to an array, in those cases + ``asarray(ar2)`` is an object array rather than the expected array of + contained values. .. versionadded:: 1.4.0 @@ -383,12 +393,13 @@ def in1d(ar1, ar2, assume_unique=False, invert=False): else: bool_ar = (sar[1:] == sar[:-1]) flag = np.concatenate((bool_ar, [invert])) - indx = order.argsort(kind='mergesort')[:len(ar1)] + ret = np.empty(ar.shape, dtype=bool) + ret[order] = flag if assume_unique: - return flag[indx] + return ret[:len(ar1)] else: - return flag[indx][rev_idx] + return ret[rev_idx] def union1d(ar1, ar2): """ @@ -417,6 +428,11 @@ def union1d(ar1, ar2): >>> np.union1d([-1, 0, 1], [-2, 0, 2]) array([-2, -1, 0, 1, 2]) + To find the union of more than two arrays, use functools.reduce: + + >>> from functools import reduce + >>> reduce(np.union1d, ([1, 3, 4, 3], [3, 1, 2, 1], [6, 3, 4, 2])) + array([1, 2, 3, 4, 6]) """ return unique(np.concatenate((ar1, ar2))) |