Pandas DataFrame分块迭代行选取Top K值实现数据处理

直接按行序分块批量处理即可，不需要逐列写循环、也不用硬编码选中行索引，处理效率远高于原apply逐列遍历的写法，365列1152行的数据可以秒级出结果。

核心参数定义 #

先统一配置规则里的固定参数，后续调整只需要改这几个值就行：

import pandas as pd
import numpy as np

BLOCK_SIZE = 48  # 每个分块的连续行数，对应1-48循环的分段规则
K = 3  # 每个分块内要选取的最大值行数
BATTERY_SIZE = 50  # 数值处理用的阈值

通用处理代码 #

以下代码默认逐列独立识别每个块内的Top K值位置，完全兼容原代码的逐列处理逻辑，哪怕同一行不同列数值差异较大也能正确运行：

# 复制原数据集，避免修改原始数据
D2j = D2i.copy()
# 按行位置生成块编号，不依赖原循环的行索引，前48行为块0，依次类推
block_labels = np.arange(len(D2i)) // BLOCK_SIZE

for _, block_data in D2j.groupby(block_labels):
    # 生成布尔掩码：标记当前块内每列值排名前K的位置，降序排名≤K即为Top K
    topk_mask = block_data.rank(ascending=False, method="min") <= K
    # 对标记位置执行自定义数值处理
    block_data[topk_mask] = block_data[topk_mask].map(
        lambda x: x - BATTERY_SIZE if x > BATTERY_SIZE else 0
    )
    # 将处理后的块写回结果集
    D2j.loc[block_data.index] = block_data

特殊场景适配 #

如果你需要按整行维度选Top K（同一行的所有列要么全处理要么全不处理，和给出的样例逻辑完全匹配），只需要把代码里生成topk_mask的部分替换成如下逻辑即可：

# 按行的均值（也可换成sum/max等聚合指标，适配实际排序规则）排序，选整行Top K
row_rank = block_data.mean(axis=1).rank(ascending=False, method="min") <= K
# 扩展为和块同形状的布尔矩阵，整行标记统一
topk_mask = np.repeat(row_rank.to_numpy().reshape(-1, 1), block_data.shape[1], axis=1)

因为提供的样例中同一行所有列的数值完全一致，两种写法跑出来的结果和给出的输出样例完全吻合。

逻辑验证说明 #

用提供的12行测试样例（把BLOCK_SIZE设为6、K=3、BATTERY_SIZE=50）测试：

• 第一个6行块内，每列Top3值对应行3、行5、行6，三个位置值分别为32、42、34，均小于50，全部置0，其余行保持原值
• 第二个6行块内，每列Top3值对应行3（54）、行4（45）、行5（43），其中54>50计算得4，45、43小于50置0，其余行保持原值
  输出结果和预期样例完全一致。

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