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 typing import (
8 TYPE_CHECKING,
9 Generic,
10 Literal,
11 cast,
12 final,
13)
15import numpy as np
17from pandas._libs import algos as libalgos
19from pandas.core.dtypes.common import (
20 is_bool_dtype,
21 is_complex_dtype,
22 is_integer_dtype,
23 is_list_like,
24 is_numeric_dtype,
25 needs_i8_conversion,
26)
27from pandas.core.dtypes.dtypes import BaseMaskedDtype
29from pandas.core.indexes.api import default_index
31if TYPE_CHECKING:
32 from collections.abc import (
33 Hashable,
34 Sequence,
35 )
37 from pandas._typing import (
38 DtypeObj,
39 IndexLabel,
40 NDFrameT,
41 )
43 from pandas import (
44 DataFrame,
45 Index,
46 Series,
47 )
48else:
49 # Generic[...] requires a non-str, provide it with a plain TypeVar at
50 # runtime to avoid circular imports
51 from pandas._typing import T
53 NDFrameT = T
54 DataFrame = T
55 Series = T
58class SelectN(Generic[NDFrameT]):
59 def __init__(
60 self, obj: NDFrameT, n: int, keep: Literal["first", "last", "all"]
61 ) -> None:
62 self.obj = obj
63 self.n = n
64 self.keep = keep
66 if self.keep not in ("first", "last", "all"):
67 raise ValueError('keep must be either "first", "last" or "all"')
69 def compute(self, method: str) -> NDFrameT:
70 raise NotImplementedError
72 @final
73 def nlargest(self) -> NDFrameT:
74 return self.compute("nlargest")
76 @final
77 def nsmallest(self) -> NDFrameT:
78 return self.compute("nsmallest")
80 @final
81 @staticmethod
82 def is_valid_dtype_n_method(dtype: DtypeObj) -> bool:
83 """
84 Helper function to determine if dtype is valid for
85 nsmallest/nlargest methods
86 """
87 if is_numeric_dtype(dtype):
88 return not is_complex_dtype(dtype)
89 return needs_i8_conversion(dtype)
92class SelectNSeries(SelectN[Series]):
93 """
94 Implement n largest/smallest for Series
96 Parameters
97 ----------
98 obj : Series
99 n : int
100 keep : {'first', 'last'}, default 'first'
102 Returns
103 -------
104 nordered : Series
105 """
107 def compute(self, method: str) -> Series:
108 from pandas.core.reshape.concat import concat
110 n = self.n
111 dtype = self.obj.dtype
112 if not self.is_valid_dtype_n_method(dtype):
113 raise TypeError(f"Cannot use method '{method}' with dtype {dtype}")
115 if n <= 0:
116 return self.obj[[]]
118 # Save index and reset to default index to avoid performance impact
119 # from when index contains duplicates
120 original_index: Index = self.obj.index
121 default_index = self.obj.reset_index(drop=True)
123 # Slower method used when taking the full length of the series
124 # In this case, it is equivalent to a sort.
125 if n >= len(default_index):
126 ascending = method == "nsmallest"
127 result = default_index.sort_values(ascending=ascending, kind="stable").head(
128 n
129 )
130 result.index = original_index.take(result.index) # type: ignore[arg-type]
131 return result
133 # Fast method used in the general case
134 dropped = default_index.dropna()
135 nan_index = default_index.drop(dropped.index)
137 new_dtype = dropped.dtype
139 # Similar to algorithms._ensure_data
140 arr = dropped._values
141 if needs_i8_conversion(arr.dtype):
142 arr = arr.view("i8")
143 elif isinstance(arr.dtype, BaseMaskedDtype):
144 arr = arr._data
145 else:
146 arr = np.asarray(arr)
147 if arr.dtype.kind == "b":
148 arr = arr.view(np.uint8)
150 if method == "nlargest":
151 arr = -arr
152 if is_integer_dtype(new_dtype):
153 # GH 21426: ensure reverse ordering at boundaries
154 arr -= 1
156 elif is_bool_dtype(new_dtype):
157 # GH 26154: ensure False is smaller than True
158 arr = 1 - (-arr)
160 if self.keep == "last":
161 arr = arr[::-1]
163 nbase = n
164 narr = len(arr)
165 n = min(n, narr)
167 # arr passed into kth_smallest must be contiguous. We copy
168 # here because kth_smallest will modify its input
169 # avoid OOB access with kth_smallest_c when n <= 0
170 if len(arr) > 0:
171 kth_val = libalgos.kth_smallest(arr.copy(order="C"), n - 1)
172 else:
173 kth_val = np.nan
174 (ns,) = np.nonzero(arr <= kth_val)
175 inds = ns[arr[ns].argsort(kind="stable")]
177 if self.keep != "all":
178 inds = inds[:n]
179 findex = nbase
180 elif len(inds) < nbase <= len(nan_index) + len(inds):
181 findex = len(nan_index) + len(inds)
182 else:
183 findex = len(inds)
185 if self.keep == "last":
186 # reverse indices
187 inds = narr - 1 - inds
189 result = concat([dropped.iloc[inds], nan_index]).iloc[:findex]
190 result.index = original_index.take(result.index)
191 return result
194class SelectNFrame(SelectN[DataFrame]):
195 """
196 Implement n largest/smallest for DataFrame
198 Parameters
199 ----------
200 obj : DataFrame
201 n : int
202 keep : {'first', 'last'}, default 'first'
203 columns : list or str
205 Returns
206 -------
207 nordered : DataFrame
208 """
210 def __init__(
211 self,
212 obj: DataFrame,
213 n: int,
214 keep: Literal["first", "last", "all"],
215 columns: IndexLabel,
216 ) -> None:
217 super().__init__(obj, n, keep)
218 if not is_list_like(columns) or isinstance(columns, tuple):
219 columns = [columns]
221 columns = cast("Sequence[Hashable]", columns)
222 columns = list(columns)
223 self.columns = columns
225 def compute(self, method: str) -> DataFrame:
226 n = self.n
227 frame = self.obj
228 columns = self.columns
230 for column in columns:
231 dtype = frame[column].dtype
232 if not self.is_valid_dtype_n_method(dtype):
233 raise TypeError(
234 f"Column {column!r} has dtype {dtype}, "
235 f"cannot use method {method!r} with this dtype"
236 )
238 def get_indexer(current_indexer: Index, other_indexer: Index) -> Index:
239 """
240 Helper function to concat `current_indexer` and `other_indexer`
241 depending on `method`
242 """
243 if method == "nsmallest":
244 return current_indexer.append(other_indexer)
245 else:
246 return other_indexer.append(current_indexer)
248 # Below we save and reset the index in case index contains duplicates
249 original_index = frame.index
250 cur_frame = frame = frame.reset_index(drop=True)
251 cur_n = n
252 indexer: Index = default_index(0)
254 for i, column in enumerate(columns):
255 # For each column we apply method to cur_frame[column].
256 # If it's the last column or if we have the number of
257 # results desired we are done.
258 # Otherwise there are duplicates of the largest/smallest
259 # value and we need to look at the rest of the columns
260 # to determine which of the rows with the largest/smallest
261 # value in the column to keep.
262 series = cur_frame[column]
263 is_last_column = len(columns) - 1 == i
264 values = getattr(series, method)(
265 cur_n, keep=self.keep if is_last_column else "all"
266 )
268 if is_last_column or len(values) <= cur_n:
269 indexer = get_indexer(indexer, values.index)
270 break
272 # Now find all values which are equal to
273 # the (nsmallest: largest)/(nlargest: smallest)
274 # from our series.
275 border_value = values == values[values.index[-1]]
277 # Some of these values are among the top-n
278 # some aren't.
279 unsafe_values = values[border_value]
281 # These values are definitely among the top-n
282 safe_values = values[~border_value]
283 indexer = get_indexer(indexer, safe_values.index)
285 # Go on and separate the unsafe_values on the remaining
286 # columns.
287 cur_frame = cur_frame.loc[unsafe_values.index]
288 cur_n = n - len(indexer)
290 frame = frame.take(indexer)
292 # Restore the index on frame
293 frame.index = original_index.take(indexer) # type: ignore[arg-type]
295 # If there is only one column, the frame is already sorted.
296 if len(columns) == 1:
297 return frame
299 ascending = method == "nsmallest"
301 return frame.sort_values(columns, ascending=ascending, kind="stable")