python常用pandas函数nlargest / nsmallest及其手动实现

import pandas as pd
 
# 创建示例DataFrame
df = pd.DataFrame({
    'A': [1, 2, 3, 4, 5],
    'B': [5, 4, 3, 2, 1],
    'C': ['a', 'b', 'c', 'd', 'e']
})
 
# 使用nlargest和nsmallest函数
# 获取每列中最大的两个值
top2_values_A = df['A'].nlargest(2)
top2_values_B = df['B'].nlargest(2)
 
# 获取每列中最小的两个值
bottom2_values_A = df['A'].nsmallest(2)
bottom2_values_B = df['B'].nsmallest(2)
 
# 手动实现nlargest和nsmallest函数
def manual_nlargest(series, n):
    return series.sort_values(ascending=False).head(n)
 
def manual_nsmallest(series, n):
    return series.sort_values(ascending=True).head(n)
 
# 手动实现的结果应该与使用pandas内置函数的结果一致
assert manual_nlargest(df['A'], 2).equals(top2_values_A)
assert manual_nsmallest(df['A'], 2).equals(bottom2_values_A)
 
print("手动实现的最大值获取成功！")

这段代码首先创建了一个简单的DataFrame，然后使用了pandas的nlargest和nsmallest函数来获取每列中最大和最小的几个值。接着，我们手动实现了这两个函数，并通过断言验证了手动实现的结果与使用内置函数的结果是否一致。这样做可以帮助理解这些函数的工作原理，并且在某些情况下可以作为替代方案，用来提高代码的可读性或者是在无法使用pandas库的环境中。