ENH: add return_inverse to df.duplicated

Author: h-vetinari

pandas/core/frame.py

@@ -4347,7 +4347,7 @@ def drop_duplicates(self, subset=None, keep='first', inplace=False):
         else:
             return self[-duplicated]
 
-    def duplicated(self, subset=None, keep='first'):
+    def duplicated(self, subset=None, keep='first', return_inverse=False):
         """
         Return boolean Series denoting duplicate rows, optionally only
         considering certain columns
@@ -4362,15 +4362,24 @@ def duplicated(self, subset=None, keep='first'):
               first occurrence.
             - ``last`` : Mark duplicates as ``True`` except for the
               last occurrence.
-            - False : Mark all duplicates as ``True``.
+            - False : Mark all duplicates as ``True``. This option is not
+              compatible with ``return_inverse``.
+        return_inverse boolean, default False
+            Determines whether the mapping from unique elements to the original
+            index should be returned. If true, the output is a tuple.
 
         Returns
         -------
-        duplicated : Series
+        duplicated : Series or tuple of Series if return_inverse is True
         """
         from pandas.core.sorting import get_group_index
         from pandas._libs.hashtable import duplicated_int64, _SIZE_HINT_LIMIT
 
+        if return_inverse and not keep:
+            raise ValueError("The parameters return_inverse=True and "
+                             "keep=False cannot be used together (impossible "
+                             "to calculate an inverse when discarding values)")
+
         def f(vals):
             labels, shape = algorithms.factorize(
                 vals, size_hint=min(len(self), _SIZE_HINT_LIMIT))
@@ -4395,7 +4404,22 @@ def f(vals):
         labels, shape = map(list, zip(*map(f, vals)))
 
         ids = get_group_index(labels, shape, sort=False, xnull=False)
-        return Series(duplicated_int64(ids, keep), index=self.index)
+        isdup = Series(duplicated_int64(ids, keep), index=self.index)
+        if not return_inverse:
+            return isdup
+
+        if keep == 'first':
+            # o2u: original indices to indices of ARRAY of unique values
+            # u2o: reduplication from array of unique values to original array
+            _, o2u, u2o = np.unique(ids, return_inverse=True,
+                                    return_index=True)
+            inv = Series(self.index[o2u][u2o], index=self.index)
+        elif keep == 'last':
+            ids = ids[::-1]  # np.unique takes first occurrence as unique value
+            _, o2u, u2o = np.unique(ids, return_inverse=True,
+                                    return_index=True)
+            inv = Series(self.index[::-1][o2u][u2o][::-1], index=self.index)
+        return isdup, inv
 
     # ----------------------------------------------------------------------
     # Sorting

pandas/tests/frame/test_duplicates.py

@@ -0,0 +1,106 @@
+# -*- coding: utf-8 -*-
+
+import pytest
+
+import numpy as np
+from pandas import Series, DataFrame
+
+from pandas.util.testing import assert_series_equal, assert_frame_equal
+import pandas.util.testing as tm
+
+
+class TestDataFrameDuplicated(object):
+
+    def test_duplicated_keep(self):
+        df = DataFrame({'A': [0, 1, 1, 2, 0], 'B': ['a', 'b', 'b', 'c', 'a']})
+
+        # keep = 'first'
+        exp = Series([False, False, True, False, True])
+        assert_series_equal(df.duplicated(keep='first'), exp)
+
+        # keep = 'last'
+        exp = Series([True, True, False, False, False])
+        assert_series_equal(df.duplicated(keep='last'), exp)
+
+        # keep = False
+        exp = Series([True, True, True, False, True])
+        assert_series_equal(df.duplicated(keep=False), exp)
+
+    def test_duplicated_nan_none(self):
+        # np.nan and None are considered equal
+        df = DataFrame({'C': [np.nan, 3, 3, None, np.nan]}, dtype=object)
+
+        # keep = 'first'
+        exp = Series([False, False, True, True, True])
+        assert_series_equal(df.duplicated(keep='first'), exp)
+
+        # keep = 'last'
+        exp = Series([True, True, False, True, False])
+        assert_series_equal(df.duplicated(keep='last'), exp)
+
+        # keep = False
+        exp = Series([True] * 5)
+        assert_series_equal(df.duplicated(keep=False), exp)
+
+    @pytest.mark.parametrize('keep', ['first', 'last', False])
+    @pytest.mark.parametrize('subset', [None, ['A', 'B'], 'A'])
+    def test_duplicated_subset(self, subset, keep):
+        df = DataFrame({'A': [0, 1, 1, 2, 0],
+                        'B': ['a', 'b', 'b', 'c', 'a'],
+                        'C': [np.nan, 3, 3, None, np.nan]})
+
+        if subset is None:
+            subset = list(df.columns)
+
+        exp = df[subset].duplicated(keep=keep).rename(name=None)
+        assert_series_equal(df.duplicated(keep=keep, subset=subset), exp)
+
+    def test_duplicated_inverse(self):
+        # check that return_inverse kwarg does not affect outcome;
+        # index of inverse must be correctly transformed as well
+        idx = [1, 4, 9, 16, 25]
+        df = DataFrame({'A': [0, 1, 1, 2, 0], 'B': ['a', 'b', 'b', 'c', 'a']},
+                       index=idx)
+
+        # keep = 'first'
+        exp_isdup = df.duplicated(keep='first')
+        exp_inv = Series([1, 4, 4, 16, 1], index=idx)
+        tst_isdup, tst_inv = df.duplicated(keep='first', return_inverse=True)
+        assert_series_equal(tst_isdup, exp_isdup)
+        assert_series_equal(tst_inv, exp_inv)
+        unique = df.loc[~exp_isdup]
+        reconstr = unique.reindex(tst_inv.values).set_index(tst_inv.index)
+        assert_frame_equal(reconstr, df)
+
+        # keep = 'last'
+        exp_isdup = df.duplicated(keep='last')
+        exp_inv = Series([25, 9, 9, 16, 25], index=idx)
+        tst_isdup, tst_inv = df.duplicated(keep='last', return_inverse=True)
+        assert_series_equal(tst_isdup, exp_isdup)
+        assert_series_equal(tst_inv, exp_inv)
+        unique = df.loc[~exp_isdup]
+        reconstr = unique.reindex(tst_inv.values).set_index(tst_inv.index)
+        assert_frame_equal(reconstr, df)
+
+        # keep = False
+        rgx = 'The parameters return_inverse=True and keep=False cannot be.*'
+        with tm.assert_raises_regex(ValueError, rgx):
+            df.duplicated(keep=False, return_inverse=True)
+
+    @pytest.mark.parametrize('keep', ['first', 'last'])
+    @pytest.mark.parametrize('subset', [None, ['A', 'B'], 'A'])
+    def test_duplicated_inverse_large(self, subset, keep):
+        # unsorted index important to check 'first'/'last' functionality
+        df = DataFrame(np.random.randint(0, 10, (10000, 3)),
+                       columns=list('ABC')).sample(5000)
+
+        exp_isdup = df.duplicated(keep=keep, subset=subset)
+        tst_isdup, inv = df.duplicated(keep=keep, subset=subset,
+                                       return_inverse=True)
+        assert_series_equal(tst_isdup, exp_isdup)
+
+        # reconstruction can only succeed if all columns are taken into account
+        if subset is None:
+            unique = df.loc[~exp_isdup]
+            reconstr = unique.reindex(inv.values).set_index(inv.index)
+            assert_frame_equal(reconstr, df)