np.argsort与pd.nsmallest取最小k值结果不一致，如何修正np.argsort用法

示例数据与基础代码 #

import pandas as pd
import numpy as np

np.random.seed(2021)
dates = pd.date_range('20130226', periods=90)
df = pd.DataFrame(np.random.uniform(0, 10, size=(90, 4)), index=dates, columns=['A_values', 'B_values', 'C_values', 'target'])

# 计算mape的函数
def mape(y_true, y_pred):
    y_pred = np.array(y_pred)
    return np.mean(np.abs(y_true - y_pred) / np.clip(np.abs(y_true), 1, np.inf),
                   axis=0)*100

preds = df.columns[df.columns.str.endswith('_values')]
k = 2
print(df)

输出结果：

A_values  B_values  C_values    target
2013-02-26  6.059783  7.333694  1.389472  3.126731
2013-02-27  9.972433  1.281624  1.789931  7.529254
2013-02-28  6.621605  7.843101  0.968944  0.585713
2013-03-01  9.623960  6.165574  0.866300  5.612724
2013-03-02  6.165247  9.638430  5.743043  3.711608
             ...       ...       ...       ...
2013-05-22  0.589729  6.479978  3.531450  6.872059
2013-05-23  6.279065  3.837670  8.853146  8.209883
2013-05-24  5.533017  5.241127  1.388056  5.355926
2013-05-25  1.596038  4.665995  2.406251  1.971875
2013-05-26  3.269001  1.787529  6.659690  7.545569

[90 rows x 4 columns]

需求说明 #

需要计算MAPE值，通过两种方法对数据按年月分组，选取每组误差最小的2个值：

方法1：错误的np.argsort实现 #

def grpProc(grp):
    err = mape(grp[preds], grp[['target']])
    print(err)
    sort_args = np.argsort(err, axis=1) < k
    cols = preds[sort_args]
    print(cols)
    print('-'*50)

df.groupby(pd.Grouper(freq='M')).apply(grpProc)

输出结果：

A_values     54.685258
B_values    212.458242
C_values    161.332752
dtype: float64
Index(['A_values', 'C_values'], dtype='object')
--------------------------------------------------
A_values     77.504315
B_values    128.986127
C_values    118.977186
dtype: float64
Index(['A_values', 'C_values'], dtype='object')
--------------------------------------------------
A_values    132.535352
B_values    150.886936
C_values     94.279492
dtype: float64
Index(['B_values', 'C_values'], dtype='object')
--------------------------------------------------
A_values    150.554314
B_values    114.113724
C_values     92.487276
dtype: float64
Index(['B_values', 'C_values'], dtype='object')
--------------------------------------------------

方法2：正确的pd.nsmallest实现 #

def grpProc(grp):
    err = mape(grp[preds], grp[['target']])
    print(err)
    cols = err.nsmallest(k).index
    print(cols)
    print('-'*50)

df.groupby(pd.Grouper(freq='M')).apply(grpProc)

输出结果：

A_values     54.685258
B_values    212.458242
C_values    161.332752
dtype: float64
Index(['A_values', 'C_values'], dtype='object')
--------------------------------------------------
A_values     77.504315
B_values    128.986127
C_values    118.977186
dtype: float64
Index(['A_values', 'C_values'], dtype='object')
--------------------------------------------------
A_values    132.535352
B_values    150.886936
C_values     94.279492
dtype: float64
Index(['C_values', 'A_values'], dtype='object')
--------------------------------------------------
A_values    150.554314
B_values    114.113724
C_values     92.487276
dtype: float64
Index(['C_values', 'B_values'], dtype='object')
--------------------------------------------------

可以看到方法1在第三组返回了错误结果，正确结果应为['C_values', 'A_values']。

正确的np.argsort实现 #

问题原因 #

原来的错误写法使用np.argsort(err, axis=1) < k生成布尔索引筛选列，会保持preds原有的列顺序，不会按误差大小排序，还会出现匹配错误。

正确代码 #

def grpProc(grp):
    err = mape(grp[preds], grp[['target']])
    print(err)
    # 获取误差升序排序后的下标，取前k个
    sorted_idx = np.argsort(err.values)[:k]
    cols = preds[sorted_idx]
    print(cols)
    print('-'*50)

df.groupby(pd.Grouper(freq='M')).apply(grpProc)

逻辑说明 #

• np.argsort对一维的误差数组默认按升序排序，返回的是原数组元素从小到大排列对应的索引下标
• 直接取前k个下标，再用preds索引这些下标，得到的列顺序和pd.nsmallest返回的顺序完全一致，都是按误差从小到大排列
• 完全复现pd.nsmallest的返回结果，不会出现匹配错误或顺序错误的问题

内容的提问来源于stack exchange，提问作者ah bon