Pandas使用sort_values实现按a列值控制b列条件排序

实现方法 #

核心思路是构造动态排序键：当a=1时用-b作为排序依据（升序排列-b等价于降序排列b），当a=-1时直接用b作为排序依据，再按time+动态键排序即可。

方法1：显式生成辅助排序列（可读性更高） #

import pandas as pd
import numpy as np

# 示例数据
df = pd.DataFrame({'time': [0, 3, 2, 2, 1], 'a': [1, -1, 1, 1, -1], 'b': [4, 5, 1, 6, 2]})

# 生成辅助排序列
df['sort_key'] = np.where(df['a'] == 1, -df['b'], df['b'])
# 排序后删除辅助列
df = df.sort_values(by=['time', 'sort_key'], ignore_index=True).drop('sort_key', axis=1)

print(df)

方法2：直接在sort_values中通过key参数实现（无需额外加列） #

import pandas as pd
import numpy as np

df = pd.DataFrame({'time': [0, 3, 2, 2, 1], 'a': [1, -1, 1, 1, -1], 'b': [4, 5, 1, 6, 2]})

df = df.sort_values(
    by=['time', 'b'],
    key=lambda col: col if col.name != 'b' else np.where(df['a'] == 1, -col, col),
    ignore_index=True
)

print(df)

两种方法的输出都和预期结果完全一致：

time  a  b
0     0  1  4
1     1 -1  2
2     2  1  6
3     2  1  1
4     3 -1  5

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