Unique1d will now return unique as well as reverse indices. Fix order of

returns [patch by Robert Cimrman].

2 files changed, 53 insertions, 12 deletions

diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py
index cd9a87931..19efbc452 100644
--- a/numpy/lib/arraysetops.py
+++ b/numpy/lib/arraysetops.py
@@ -74,7 +74,7 @@ def ediff1d(ary, to_end=None, to_begin=None):
 
     return ed
 
-def unique1d(ar1, return_index=False):
+def unique1d(ar1, return_index=False, return_inverse=False):
     """
     Find the unique elements of an array.
 
@@ -85,6 +85,9 @@ def unique1d(ar1, return_index=False):
     return_index : bool, optional
         If True, also return the indices against `ar1` that result in the
         unique array.
+    return_inverse : bool, optional
+        If True, also return the indices against the unique array that
+        result in `ar1`.
 
     Returns
     -------
@@ -93,6 +96,9 @@ def unique1d(ar1, return_index=False):
     unique_indices : ndarray, optional
         The indices of the unique values. Only provided if `return_index` is
         True.
+    unique_inverse : ndarray, optional
+        The indices to reconstruct the original array. Only provided if
+        `return_inverse` is True.
 
     See Also
     --------
@@ -107,21 +113,52 @@ def unique1d(ar1, return_index=False):
     >>> np.unique1d(a)
     array([1, 2, 3])
 
+    Reconstruct the input from unique values:
+
+    >>> np.unique1d([1,2,6,4,2,3,2], return_index=True)
+    >>> x = [1,2,6,4,2,3,2]
+    >>> u, i = np.unique1d(x, return_inverse=True)
+    >>> u
+    array([1, 2, 3, 4, 6])
+    >>> i
+    array([0, 1, 4, 3, 1, 2, 1])
+    >>> [u[p] for p in i]
+    [1, 2, 6, 4, 2, 3, 2]
+
     """
+    if return_index:
+        import warnings
+        warnings.warn("The order of the output arguments for "
+                      "`return_index` has changed.  Before, "
+                      "the output was (indices, unique_arr), but "
+                      "has now been reversed to be more consistent.")
+
     ar = np.asarray(ar1).flatten()
     if ar.size == 0:
-        if return_index: return np.empty(0, np.bool), ar
-        else: return ar
-
-    if return_index:
+        if return_inverse and return_index:
+            return ar, np.empty(0, np.bool), np.empty(0, np.bool)
+        elif return_inverse or return_index:
+            return ar, np.empty(0, np.bool)
+        else:
+            return ar
+
+    if return_inverse or return_index:
         perm = ar.argsort()
         aux = ar[perm]
-        flag = np.concatenate( ([True], aux[1:] != aux[:-1]) )
-        return perm[flag], aux[flag]
+        flag = np.concatenate(([True], aux[1:] != aux[:-1]))
+        if return_inverse:
+            iflag = np.cumsum(flag) - 1
+            iperm = perm.argsort()
+            if return_index:
+                return aux[flag], perm[flag], iflag[iperm]
+            else:
+                return aux[flag], iflag[iperm]
+        else:
+            return aux[flag], perm[flag]
 
     else:
         ar.sort()
-        flag = np.concatenate( ([True], ar[1:] != ar[:-1]) )
+        flag = np.concatenate(([True], ar[1:] != ar[:-1]))
         return ar[flag]
 
 def intersect1d(ar1, ar2):
diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py
index 474842c8a..c40cf9d20 100644
--- a/numpy/lib/tests/test_arraysetops.py
+++ b/numpy/lib/tests/test_arraysetops.py
@@ -6,6 +6,8 @@ from numpy.testing import *
 import numpy as np
 from numpy.lib.arraysetops import *
 
+import warnings
+
 class TestAso(TestCase):
     def test_unique1d( self ):
         a = np.array( [5, 7, 1, 2, 1, 5, 7] )
@@ -14,11 +16,13 @@ class TestAso(TestCase):
         c = unique1d( a )
         assert_array_equal( c, ec )
 
-        d, c = unique1d( a, True )
-        ed = np.array( [2, 3, 0, 1] )
+        warnings.simplefilter('ignore', Warning)
+        unique, indices = unique1d( a, return_index=True )
+        warnings.resetwarnings()
 
-        assert_array_equal( d,ed )
-        assert_array_equal( c, ec )
+        ed = np.array( [2, 3, 0, 1] )
+        assert_array_equal(unique, ec)
+        assert_array_equal(indices, ed)
 
         assert_array_equal([], unique1d([]))