Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Implementation of nlargest and nsmallest.
3"""
5from __future__ import annotations
7from collections.abc import (
8 Hashable,
9 Sequence,
10)
11from typing import (
12 TYPE_CHECKING,
13 Generic,
14 Literal,
15 cast,
16 final,
17)
19import numpy as np
21from pandas._libs import algos as libalgos
23from pandas.core.dtypes.common import (
24 is_bool_dtype,
25 is_complex_dtype,
26 is_integer_dtype,
27 is_list_like,
28 is_numeric_dtype,
29 needs_i8_conversion,
30)
31from pandas.core.dtypes.dtypes import BaseMaskedDtype
33from pandas.core.indexes.api import default_index
35if TYPE_CHECKING:
36 from pandas._typing import (
37 DtypeObj,
38 IndexLabel,
39 NDFrameT,
40 )
42 from pandas import (
43 DataFrame,
44 Index,
45 Series,
46 )
47else:
48 # Generic[...] requires a non-str, provide it with a plain TypeVar at
49 # runtime to avoid circular imports
50 from pandas._typing import T
52 NDFrameT = T
53 DataFrame = T
54 Series = T
57class SelectN(Generic[NDFrameT]):
58 def __init__(
59 self, obj: NDFrameT, n: int, keep: Literal["first", "last", "all"]
60 ) -> None:
61 self.obj = obj
62 self.n = n
63 self.keep = keep
65 if self.keep not in ("first", "last", "all"):
66 raise ValueError('keep must be either "first", "last" or "all"')
68 def compute(self, method: str) -> NDFrameT:
69 raise NotImplementedError
71 @final
72 def nlargest(self) -> NDFrameT:
73 return self.compute("nlargest")
75 @final
76 def nsmallest(self) -> NDFrameT:
77 return self.compute("nsmallest")
79 @final
80 @staticmethod
81 def is_valid_dtype_n_method(dtype: DtypeObj) -> bool:
82 """
83 Helper function to determine if dtype is valid for
84 nsmallest/nlargest methods
85 """
86 if is_numeric_dtype(dtype):
87 return not is_complex_dtype(dtype)
88 return needs_i8_conversion(dtype)
91class SelectNSeries(SelectN[Series]):
92 """
93 Implement n largest/smallest for Series
95 Parameters
96 ----------
97 obj : Series
98 n : int
99 keep : {'first', 'last'}, default 'first'
101 Returns
102 -------
103 nordered : Series
104 """
106 def compute(self, method: str) -> Series:
107 from pandas.core.reshape.concat import concat
109 n = self.n
110 dtype = self.obj.dtype
111 if not self.is_valid_dtype_n_method(dtype):
112 raise TypeError(f"Cannot use method '{method}' with dtype {dtype}")
114 if n <= 0:
115 return self.obj[[]]
117 # Save index and reset to default index to avoid performance impact
118 # from when index contains duplicates
119 original_index: Index = self.obj.index
120 default_index = self.obj.reset_index(drop=True)
122 # Slower method used when taking the full length of the series
123 # In this case, it is equivalent to a sort.
124 if n >= len(default_index):
125 ascending = method == "nsmallest"
126 result = default_index.sort_values(ascending=ascending, kind="stable").head(
127 n
128 )
129 result.index = original_index.take(result.index)
130 return result
132 # Fast method used in the general case
133 dropped = default_index.dropna()
134 nan_index = default_index.drop(dropped.index)
136 new_dtype = dropped.dtype
138 # Similar to algorithms._ensure_data
139 arr = dropped._values
140 if needs_i8_conversion(arr.dtype):
141 arr = arr.view("i8")
142 elif isinstance(arr.dtype, BaseMaskedDtype):
143 arr = arr._data
144 else:
145 arr = np.asarray(arr)
146 if arr.dtype.kind == "b":
147 arr = arr.view(np.uint8)
149 if method == "nlargest":
150 arr = -arr
151 if is_integer_dtype(new_dtype):
152 # GH 21426: ensure reverse ordering at boundaries
153 arr -= 1
155 elif is_bool_dtype(new_dtype):
156 # GH 26154: ensure False is smaller than True
157 arr = 1 - (-arr)
159 if self.keep == "last":
160 arr = arr[::-1]
162 nbase = n
163 narr = len(arr)
164 n = min(n, narr)
166 # arr passed into kth_smallest must be contiguous. We copy
167 # here because kth_smallest will modify its input
168 # avoid OOB access with kth_smallest_c when n <= 0
169 if len(arr) > 0:
170 kth_val = libalgos.kth_smallest(arr.copy(order="C"), n - 1)
171 else:
172 kth_val = np.nan
173 (ns,) = np.nonzero(arr <= kth_val)
174 inds = ns[arr[ns].argsort(kind="stable")]
176 if self.keep != "all":
177 inds = inds[:n]
178 findex = nbase
179 elif len(inds) < nbase <= len(nan_index) + len(inds):
180 findex = len(nan_index) + len(inds)
181 else:
182 findex = len(inds)
184 if self.keep == "last":
185 # reverse indices
186 inds = narr - 1 - inds
188 result = concat([dropped.iloc[inds], nan_index]).iloc[:findex]
189 result.index = original_index.take(result.index)
190 return result
193class SelectNFrame(SelectN[DataFrame]):
194 """
195 Implement n largest/smallest for DataFrame
197 Parameters
198 ----------
199 obj : DataFrame
200 n : int
201 keep : {'first', 'last'}, default 'first'
202 columns : list or str
204 Returns
205 -------
206 nordered : DataFrame
207 """
209 def __init__(
210 self,
211 obj: DataFrame,
212 n: int,
213 keep: Literal["first", "last", "all"],
214 columns: IndexLabel,
215 ) -> None:
216 super().__init__(obj, n, keep)
217 if not is_list_like(columns) or isinstance(columns, tuple):
218 columns = [columns]
220 columns = cast(Sequence[Hashable], columns)
221 columns = list(columns)
222 self.columns = columns
224 def compute(self, method: str) -> DataFrame:
225 n = self.n
226 frame = self.obj
227 columns = self.columns
229 for column in columns:
230 dtype = frame[column].dtype
231 if not self.is_valid_dtype_n_method(dtype):
232 raise TypeError(
233 f"Column {column!r} has dtype {dtype}, "
234 f"cannot use method {method!r} with this dtype"
235 )
237 def get_indexer(current_indexer: Index, other_indexer: Index) -> Index:
238 """
239 Helper function to concat `current_indexer` and `other_indexer`
240 depending on `method`
241 """
242 if method == "nsmallest":
243 return current_indexer.append(other_indexer)
244 else:
245 return other_indexer.append(current_indexer)
247 # Below we save and reset the index in case index contains duplicates
248 original_index = frame.index
249 cur_frame = frame = frame.reset_index(drop=True)
250 cur_n = n
251 indexer: Index = default_index(0)
253 for i, column in enumerate(columns):
254 # For each column we apply method to cur_frame[column].
255 # If it's the last column or if we have the number of
256 # results desired we are done.
257 # Otherwise there are duplicates of the largest/smallest
258 # value and we need to look at the rest of the columns
259 # to determine which of the rows with the largest/smallest
260 # value in the column to keep.
261 series = cur_frame[column]
262 is_last_column = len(columns) - 1 == i
263 values = getattr(series, method)(
264 cur_n, keep=self.keep if is_last_column else "all"
265 )
267 if is_last_column or len(values) <= cur_n:
268 indexer = get_indexer(indexer, values.index)
269 break
271 # Now find all values which are equal to
272 # the (nsmallest: largest)/(nlargest: smallest)
273 # from our series.
274 border_value = values == values[values.index[-1]]
276 # Some of these values are among the top-n
277 # some aren't.
278 unsafe_values = values[border_value]
280 # These values are definitely among the top-n
281 safe_values = values[~border_value]
282 indexer = get_indexer(indexer, safe_values.index)
284 # Go on and separate the unsafe_values on the remaining
285 # columns.
286 cur_frame = cur_frame.loc[unsafe_values.index]
287 cur_n = n - len(indexer)
289 frame = frame.take(indexer)
291 # Restore the index on frame
292 frame.index = original_index.take(indexer)
294 # If there is only one column, the frame is already sorted.
295 if len(columns) == 1:
296 return frame
298 ascending = method == "nsmallest"
300 return frame.sort_values(columns, ascending=ascending, kind="stable")