diff options
Diffstat (limited to 'numpy/lib/arraysetops.py')
| -rw-r--r-- | numpy/lib/arraysetops.py | 70 |
1 files changed, 44 insertions, 26 deletions
diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py index d3b6119f4..836f4583f 100644 --- a/numpy/lib/arraysetops.py +++ b/numpy/lib/arraysetops.py @@ -78,29 +78,46 @@ def ediff1d(ary, to_end=None, to_begin=None): array([ 1, 2, -3, 5, 18]) """ - ary = np.asanyarray(ary).flat - ed = ary[1:] - ary[:-1] - arrays = [ed] - if to_begin is not None: - arrays.insert(0, to_begin) - if to_end is not None: - arrays.append(to_end) + # force a 1d array + ary = np.asanyarray(ary).ravel() - if len(arrays) != 1: - # We'll save ourselves a copy of a potentially large array in - # the common case where neither to_begin or to_end was given. - ed = np.hstack(arrays) + # fast track default case + if to_begin is None and to_end is None: + return ary[1:] - ary[:-1] + + if to_begin is None: + l_begin = 0 + else: + to_begin = np.asanyarray(to_begin).ravel() + l_begin = len(to_begin) + + if to_end is None: + l_end = 0 + else: + to_end = np.asanyarray(to_end).ravel() + l_end = len(to_end) + + # do the calculation in place and copy to_begin and to_end + l_diff = max(len(ary) - 1, 0) + result = np.empty(l_diff + l_begin + l_end, dtype=ary.dtype) + result = ary.__array_wrap__(result) + if l_begin > 0: + result[:l_begin] = to_begin + if l_end > 0: + result[l_begin + l_diff:] = to_end + np.subtract(ary[1:], ary[:-1], result[l_begin:l_begin + l_diff]) + return result - return ed def unique(ar, return_index=False, return_inverse=False, return_counts=False): """ Find the unique elements of an array. - Returns the sorted unique elements of an array. There are two optional + Returns the sorted unique elements of an array. There are three optional outputs in addition to the unique elements: the indices of the input array - that give the unique values, and the indices of the unique array that - reconstruct the input array. + that give the unique values, the indices of the unique array that + reconstruct the input array, and the number of times each unique value + comes up in the input array. Parameters ---------- @@ -113,10 +130,11 @@ def unique(ar, return_index=False, return_inverse=False, return_counts=False): If True, also return the indices of the unique array that can be used to reconstruct `ar`. return_counts : bool, optional - .. versionadded:: 1.9.0 If True, also return the number of times each unique value comes up in `ar`. + .. versionadded:: 1.9.0 + Returns ------- unique : ndarray @@ -128,10 +146,11 @@ def unique(ar, return_index=False, return_inverse=False, return_counts=False): The indices to reconstruct the (flattened) original array from the unique array. Only provided if `return_inverse` is True. unique_counts : ndarray, optional - .. versionadded:: 1.9.0 The number of times each of the unique values comes up in the original array. Only provided if `return_counts` is True. + .. versionadded:: 1.9.0 + See Also -------- numpy.lib.arraysetops : Module with a number of other functions for @@ -392,12 +411,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): """ @@ -468,11 +488,9 @@ def setdiff1d(ar1, ar2, assume_unique=False): array([1, 2]) """ - if not assume_unique: + if assume_unique: + ar1 = np.asarray(ar1).ravel() + else: ar1 = unique(ar1) ar2 = unique(ar2) - aux = in1d(ar1, ar2, assume_unique=True) - if aux.size == 0: - return aux - else: - return np.asarray(ar1)[aux == 0] + return ar1[in1d(ar1, ar2, assume_unique=True, invert=True)] |