Fix nsmallest/nlargest With Identical Values

Author: Roger Thomas

doc/source/whatsnew/v0.20.0.txt

@@ -1057,3 +1057,4 @@ Bug Fixes
 - Bug in ``pd.melt()`` where passing a tuple value for ``value_vars`` caused a ``TypeError`` (:issue:`15348`)
 - Bug in ``.eval()`` which caused multiline evals to fail with local variables not on the first line (:issue:`15342`)
 - Bug in ``pd.read_msgpack()`` which did not allow to load dataframe with an index of type ``CategoricalIndex`` (:issue:`15487`)
+- Bug in ``DataFrame.nsmallest`` and ``DataFrame.nlargest`` where identical values resulted in duplicated rows (:issue:`15297`)

pandas/core/algorithms.py

@@ -944,14 +944,27 @@ def select_n_frame(frame, columns, n, method, keep):
     -------
     nordered : DataFrame
     """
-    from pandas.core.series import Series
     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]
+    else:
+        columns = list(columns)
+    reverse = method == 'nlargest'
+    for i, column in enumerate(columns):
+        series = frame[column]
+        if reverse:
+            inds = series.argsort()[::-1][:n]
+        else:
+            inds = series.argsort()[:n]
+        values = series.take(inds)
+        if i != len(columns) - 1 and values.duplicated().any():
+            # This series has duplicate values => we must consider all rows in
+            # frame that match `values`
+            # The first condition is for the last column. In this case we don't
+            # care if there are duplicates => no need to do the check
+            frame = frame[series.isin(values)]
+        else:
+            break
+    return frame.take(inds)
 
 
 def _finalize_nsmallest(arr, kth_val, n, keep, narr):

pandas/tests/frame/test_analytics.py

@@ -1140,6 +1140,18 @@ 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': [5, 4, 3]}, index=[4, 3, 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 +1171,32 @@ 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)
-        tm.assert_frame_equal(result, expected)
+        df = pd.DataFrame({'a': [1, 2, 3, 3, 3],
+                           'b': [1, 1, 1, 1, 1],
+                           'c': [0, 1, 2, 5, 4]},
+                          index=[0, 0, 1, 1, 1])
 
-        result = df.nlargest(4, 'a')
-        expected = df.sort_values('a', ascending=False).head(4)
+        result = df.nsmallest(4, ['a', 'b', 'c'])
+        expected = df.sort_values(['a', 'b', 'c']).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, ['a', 'b', 'c'])
+        expected = df.sort_values(['a', 'b', 'c'], 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.nlargest(4, ['c', 'b', 'a'])
+        expected = df.sort_values(['c', 'b', 'a'], ascending=False).head(4)
         tm.assert_frame_equal(result, expected)
 
-        result = df.nlargest(4, ['a', 'c'])
-        expected = df.sort_values(['a', 'c'], ascending=False).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, ['c', 'a'])
-        expected = df.sort_values(['c', 'a'], 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)
+
     # ----------------------------------------------------------------------
     # Isin
 

pandas/tests/series/test_analytics.py

@@ -1455,6 +1455,16 @@ 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 = Series([1] * 3, index=[1, 2, 3])
+
+        result = s.nsmallest(3)
+        assert_series_equal(result, expected)
+
+        result = s.nlargest(3)
+        assert_series_equal(result, expected)
+
     def test_sort_index_level(self):
         mi = MultiIndex.from_tuples([[1, 1, 3], [1, 1, 1]], names=list('ABC'))
         s = Series([1, 2], mi)