Pandas中按组排序并根据组聚合值对组进行排序的技巧

import pandas as pd
import numpy as np

df = pd.DataFrame({'col1': ['A', 'B', 'A', 'B', 'C'],
                   'col2': [3, 1, 2, 4, 3],
                   'col3': [10, 20, 30, 40, 50]})
print("原始DataFrame:")
print(df)

原始DataFrame:
  col1  col2  col3
0    A     3    10
1    B     1    20
2    A     2    30
3    B     4    40
4    C     3    50

  col1  col2  col3
1    B     1    20
3    B     4    40
0    A     3    10
2    A     2    30
4    C     3    50

print("\ndf.sort_values(['col1', 'col2']):")
print(df.sort_values(['col1', 'col2']))

df.sort_values(['col1', 'col2']):
  col1  col2  col3
2    A     2    30
0    A     3    10
1    B     1    20
3    B     4    40
4    C     3    50

df_temp = df.copy()
df_temp['min_col2'] = df_temp.groupby('col1')['col2'].transform('min')
sorted_df_temp = df_temp.sort_values(['min_col2', 'col1', 'col2']).drop("min_col2", axis="columns")
print("\n使用临时列的方法:")
print(sorted_df_temp)

# 规范解决方案
# 1. 计算每个组的排序依据值（例如，每个col1组的col2最小值）
group_min_col2 = df.groupby('col1')['col2'].transform('min')

# 2. 获取这些值的排序索引
# np.argsort返回的是排序后的元素在原始数组中的位置索引
sorted_indices = np.argsort(group_min_col2)

# 3. 使用iloc根据这些索引重新排列DataFrame
out_df = df.iloc[sorted_indices]

# 为了同时实现组内排序，可以先进行一次常规排序，再进行组间排序
# 或者在iloc之后，对每个组进行内部排序
# 更简洁的方式是，在argsort之前，确保数据已经按照组内规则排序
# 最佳实践是，先对df进行一次常规的按组和组内列的排序，然后使用argsort来调整组的顺序
# 这里的需求是：B组1，4；A组2，3；C组3。然后组之间B<A<C。
# 所以，我们希望先按col1, col2排序，然后调整组的顺序。
# 实际上，上面的argsort已经给出了组的顺序，但组内没有排序。
# 要达到最终目标（组间按min_col2排序，组内按col2排序），需要结合多个排序步骤。

# 正确的组合方式：
# 1. 首先，对DataFrame进行一次常规的组内排序
df_sorted_within_groups = df.sort_values(['col1', 'col2'])

# 2. 然后，基于原始DataFrame计算用于组间排序的依据
group_min_col2_original_df = df.groupby('col1')['col2'].transform('min')

# 3. 获取基于组间排序依据的索引
# 注意：这里argsort应该作用于原始DataFrame的索引，以确保我们能正确重排
# 更直接的方法是，argsort作用于一个Series，这个Series的索引与原始DataFrame的索引一致
# 我们可以创建一个包含排序键和原始索引的DataFrame，然后对其进行排序
# 这里的挑战在于，我们希望的输出是先B组，然后A组，然后C组，并且B组内部是1,4，A组内部是2,3。
# 原始的argsort方法可以做到组的整体排序，但无法保证组内的顺序。

# 重新审视期望输出：
#   col1 col2    col3
# 1   B   1   20  (df index 1)
# 3   B   4   40  (df index 3)
# 0   A   3   10  (df index 0)
# 2   A   2   30  (df index 2)
# 4   C   3   50  (df index 4)

# 这里的关键是，我们首先需要知道每个原始行属于哪个“组排序级别”。
# B组的min_col2是1，A组是2，C组是3。
# 我们可以创建一个新的排序键，包含组的min_col2和行本身的col2。

# 步骤1：计算每个行所属组的最小值
df['group_min_col2'] = df.groupby('col1')['col2'].transform('min')

# 步骤2：根据新的排序键进行多级排序
# 优先按组的最小值排序，然后按col1（确保同一组的行在一起），最后按col2进行组内排序
final_out = df.sort_values(by=['group_min_col2', 'col1', 'col2']).drop(columns=['group_min_col2'])

print("\n规范解决方案 (多级排序):")
print(final_out)

# 原始答案中的解决方案
# 核心思想是：生成一个与df行数相同的Series，其值代表了行所属组的排序优先级。
# 然后对这个Series进行argsort，得到最终的行索引顺序。
# 这里的挑战在于，如何让argsort同时考虑组间排序和组内排序。
# 答案中提供的out = df.iloc[np.argsort(df.groupby('col1')['col2'].transform('min'))]
# 这个代码片段只保证了组的顺序，但组内的顺序是原始的，不是按col2排序的。

# 为了实现“组间按min_col2排序，组内按col2排序”：
# 我们可以创建一个复合的排序键，然后对这个键进行argsort。
# 复合键的思路是：将组的排序值（如min_col2）和行自身的col2值组合起来。
# 例如，可以创建一个元组列表：[(min_col2_for_row_i, col2_for_row_i), ...]

# 步骤1：计算每个行所属组的最小值
min_col2_series = df.groupby('col1')['col2'].transform('min')

# 步骤2：创建复合排序键
# 注意：这里需要确保argsort作用于一个能够反映最终排序顺序的单一序列。
# 将 min_col2_series 和 df['col2'] 结合起来，并对它们进行argsort。
# np.lexsort 可以用于多列排序的索引。
# lexsort(keys, axis=-1): Perform an indirect stable sort using a sequence of keys.
# keys: (k_n, k_n-1, ..., k_0) - keys are sorted from last to first.
# 所以，我们想要先按 min_col2_series 排序，再按 df['col2'] 排序。
# lexsort 的 keys 顺序是：最后排序的键在最前面，最先排序的键在最后面。
# 也就是说，如果想先按A排，再按B排，那么keys=(B, A)。
# 所以，我们想先按 min_col2_series 排，再按 df['col2'] 排，那么 keys=(df['col2'], min_col2_series)。
sorted_indices_complex = np.lexsort((df['col2'], min_col2_series))

final_out_lexsort = df.iloc[sorted_indices_complex]

print("\n使用 np.lexsort 的规范解决方案:")
print(final_out_lexsort)

out_pipeline = df.iloc[lambda d: np.lexsort((d['col2'], d.groupby('col1')['col2'].transform('min')))]
print("\n管道中使用的 np.lexsort 解决方案:")
print(out_pipeline)

out_key = df.sort_values(by='col2',
                         key=lambda s: s.groupby(df['col1']).transform('min'))
print("\n使用 sort_values 的 key 参数解决方案:")
print(out_key)