Skip to content

Commit 77ab2e6

Browse files
committed
address comments
1 parent de870a6 commit 77ab2e6

File tree

1 file changed

+16
-15
lines changed

1 file changed

+16
-15
lines changed

pandas/core/methods/selectn.py

+16-15
Original file line numberDiff line numberDiff line change
@@ -117,21 +117,22 @@ def compute(self, method: str) -> Series:
117117
# Save index and reset to default index to avoid performance impact
118118
# from when index contains duplicates
119119
original_index: Index = self.obj.index
120-
cur_series = self.obj.reset_index(drop=True)
120+
default_index = self.obj.reset_index(drop=True)
121121

122-
# slow method
123-
if n >= len(cur_series):
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):
124125
ascending = method == "nsmallest"
125-
final_series = cur_series.sort_values(
126-
ascending=ascending, kind="stable"
127-
).head(n)
128-
final_series.index = original_index.take(final_series.index)
129-
return final_series
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
130131

131-
dropped = cur_series.dropna()
132-
nan_index = cur_series.drop(dropped.index)
132+
# Fast method used in the general case
133+
dropped = default_index.dropna()
134+
nan_index = default_index.drop(dropped.index)
133135

134-
# fast method
135136
new_dtype = dropped.dtype
136137

137138
# Similar to algorithms._ensure_data
@@ -170,7 +171,7 @@ def compute(self, method: str) -> Series:
170171
else:
171172
kth_val = np.nan
172173
(ns,) = np.nonzero(arr <= kth_val)
173-
inds = ns[arr[ns].argsort(kind="mergesort")]
174+
inds = ns[arr[ns].argsort(kind="stable")]
174175

175176
if self.keep != "all":
176177
inds = inds[:n]
@@ -185,9 +186,9 @@ def compute(self, method: str) -> Series:
185186
# reverse indices
186187
inds = narr - 1 - inds
187188

188-
final_series = concat([dropped.iloc[inds], nan_index]).iloc[:findex]
189-
final_series.index = original_index.take(final_series.index)
190-
return final_series
189+
result = concat([dropped.iloc[inds], nan_index]).iloc[:findex]
190+
result.index = original_index.take(result.index)
191+
return result
191192

192193

193194
class SelectNFrame(SelectN[DataFrame]):

0 commit comments

Comments
 (0)