Fix nsmallest/nlargest With Identical Values

Remove

Add comments

Author: Roger Thomas

doc/source/whatsnew/v0.20.0.txt

@@ -1031,6 +1031,7 @@ Reshaping
 - Bug in ``pd.pivot_table()`` where no error was raised when values argument was not in the columns (:issue:`14938`)
 - Bug in ``pd.concat()`` in which concatting with an empty dataframe with ``join='inner'`` was being improperly handled (:issue:`15328`)
 - Bug with ``sort=True`` in ``DataFrame.join`` and ``pd.merge`` when joining on indexes (:issue:`15582`)
+- Bug in ``DataFrame.nsmallest`` and ``DataFrame.nlargest`` where identical values resulted in duplicated rows (:issue:`15297`)
 
 Numeric
 ^^^^^^^

pandas/core/algorithms.py

@@ -948,10 +948,35 @@ def select_n_frame(frame, columns, n, method, keep):
     if not is_list_like(columns):
         columns = [columns]
     columns = list(columns)
-    ser = getattr(frame[columns[0]], method)(n, keep=keep)
-    if isinstance(ser, Series):
-        ser = ser.to_frame()
-    return ser.merge(frame, on=columns[0], left_index=True)[frame.columns]
+    ascending = method == 'nsmallest'
+    original_frame, original_index = frame, frame.index
+    frame.reset_index(drop=True, inplace=True)
+    for i, column in enumerate(columns):
+        # For each column in columns we peform ``method`` on this frame
+        # To guard against the possibility ``method`` column has duplicate
+        # values that must be considered for futher columns (# GH15297) we
+        # filter using isin on the values returned by ``method``. If there are
+        # no duplicated values, we simply reindex like the values returned
+        # by ``method``, otherwise we sort the frame and continue
+        series = frame[column]
+        values = getattr(series, method)(n, keep=keep)
+        if i + 1 == len(columns):
+            # This is the last column => duplicates here don't matter
+            frame = frame.reindex(values.index)
+        else:
+            filtered_frame = frame[series.isin(values)]
+            if len(filtered_frame) == len(values):
+                # Values are unique in series => reindex and break
+                frame = frame.reindex(values.index)
+                break
+            # Values are not unique in series => sort and continue
+            frame = filtered_frame.sort_values(
+                column, ascending=ascending
+            )
+    original_frame.index = original_index  # Restore the index
+    # Below we set the index of the returning frame to the original index
+    frame.index = original_index[frame.index]
+    return frame
 
 
 def _finalize_nsmallest(arr, kth_val, n, keep, narr):

pandas/tests/frame/test_analytics.py

@@ -1140,6 +1140,20 @@ def test_nlargest_multiple_columns(self):
         expected = df.sort_values(['a', 'b'], ascending=False).head(5)
         tm.assert_frame_equal(result, expected)
 
+    def test_nlargest_nsmallest_identical_values(self):
+        # GH15297
+        df = pd.DataFrame({'a': [1] * 5, 'b': [1, 2, 3, 4, 5]})
+
+        result = df.nlargest(3, 'a')
+        expected = pd.DataFrame(
+            {'a': [1] * 3, 'b': [1, 2, 3]}, index=[0, 1, 2]
+        )
+        tm.assert_frame_equal(result, expected)
+
+        result = df.nsmallest(3, 'a')
+        expected = pd.DataFrame({'a': [1] * 3, 'b': [1, 2, 3]})
+        tm.assert_frame_equal(result, expected)
+
     def test_nsmallest(self):
         from string import ascii_lowercase
         df = pd.DataFrame({'a': np.random.permutation(10),
@@ -1159,33 +1173,41 @@ def test_nsmallest_multiple_columns(self):
 
     def test_nsmallest_nlargest_duplicate_index(self):
         # GH 13412
-        df = pd.DataFrame({'a': [1, 2, 3, 4],
-                           'b': [4, 3, 2, 1],
-                           'c': [0, 1, 2, 3]},
-                          index=[0, 0, 1, 1])
-        result = df.nsmallest(4, 'a')
-        expected = df.sort_values('a').head(4)
+        df = pd.DataFrame({'a': [1, 2, 3, 4, 4],
+                           'b': [1, 1, 1, 1, 1],
+                           'c': [0, 1, 2, 5, 4]},
+                          index=[0, 0, 1, 1, 1])
+
+        result = df.nsmallest(4, ['a', 'b', 'c'])
+        expected = df.sort_values(['a', 'b', 'c']).head(4)
         tm.assert_frame_equal(result, expected)
 
-        result = df.nlargest(4, 'a')
-        expected = df.sort_values('a', ascending=False).head(4)
+        result = df.nlargest(4, ['a', 'b', 'c'])
+        expected = df.sort_values(['a', 'b', 'c'], ascending=False).head(4)
         tm.assert_frame_equal(result, expected)
 
-        result = df.nsmallest(4, ['a', 'c'])
-        expected = df.sort_values(['a', 'c']).head(4)
+        result = df.nlargest(4, ['c', 'b', 'a'])
+        expected = df.sort_values(['c', 'b', 'a'], ascending=False).head(4)
         tm.assert_frame_equal(result, expected)
 
-        result = df.nsmallest(4, ['c', 'a'])
-        expected = df.sort_values(['c', 'a']).head(4)
+        result = df.nsmallest(4, ['c', 'b', 'a'])
+        expected = df.sort_values(['c', 'b', 'a']).head(4)
         tm.assert_frame_equal(result, expected)
 
-        result = df.nlargest(4, ['a', 'c'])
-        expected = df.sort_values(['a', 'c'], ascending=False).head(4)
+        # Test all duplicates still returns df of size n
+        result = df.nsmallest(2, 'b')
+        expected = df.sort_values('b').head(2)
         tm.assert_frame_equal(result, expected)
 
-        result = df.nlargest(4, ['c', 'a'])
-        expected = df.sort_values(['c', 'a'], ascending=False).head(4)
+    def test_nsmallest_nlargest_duplicate_multi_index(self):
+        df = pd.DataFrame({'a': [1, 2, 3, 3, 3],
+                           'b': [1, 1, 1, 1, 1],
+                           'c': [0, 1, 2, 5, 4]},
+                          index=[[0, 0, 0, 0, 0], [1, 1, 1, 1, 1]])
+        result = df.nsmallest(4, ['a', 'b', 'c'])
+        expected = df.sort_values(['a', 'b', 'c']).head(4)
         tm.assert_frame_equal(result, expected)
+
     # ----------------------------------------------------------------------
     # Isin
 

pandas/tests/series/test_analytics.py

@@ -1455,6 +1455,23 @@ def test_nsmallest_nlargest(self):
         expected = s.sort_values().head(3)
         assert_series_equal(result, expected)
 
+        # GH 15297
+        s = Series([1] * 5, index=[1, 2, 3, 4, 5])
+        expected_first = Series([1] * 3, index=[1, 2, 3])
+        expected_last = Series([1] * 3, index=[5, 4, 3])
+
+        result = s.nsmallest(3)
+        assert_series_equal(result, expected_first)
+
+        result = s.nsmallest(3, keep='last')
+        assert_series_equal(result, expected_last)
+
+        result = s.nlargest(3)
+        assert_series_equal(result, expected_first)
+
+        result = s.nlargest(3, keep='last')
+        assert_series_equal(result, expected_last)
+
     def test_sort_index_level(self):
         mi = MultiIndex.from_tuples([[1, 1, 3], [1, 1, 1]], names=list('ABC'))
         s = Series([1, 2], mi)