如何用Numpy按ids数组唯一索引对dst行元素分组求和并取最大值

实现思路 #

• 采用全向量化操作避免Python级循环，核心使用np.bincount完成分组统计，性能远高于逐行循环实现，数组规模越大优势越明显
• 为不同行的id添加全局唯一偏移量，避免不同行的同值id互相干扰导致统计错误
• 单次调用np.bincount即可同时完成全量分组计数、分组求和，最后按行提取最大值对应的结果即可

实现代码 #

import numpy as np

def calc_max_group(dst: np.ndarray, ids: np.ndarray) -> np.ndarray:
    # 输入维度校验
    assert dst.shape == ids.shape, "dst和ids的行列数必须完全一致"
    n_rows, n_cols = dst.shape
    
    # 计算偏移量保证不同行的id全局唯一
    id_base = ids.max() + 1
    row_offsets = np.arange(n_rows) * id_base
    global_unique_ids = (ids + row_offsets[:, None]).ravel()
    
    # 批量统计分组计数、分组求和结果
    group_counts = np.bincount(global_unique_ids)
    group_sums = np.bincount(global_unique_ids, weights=dst.ravel())
    
    # 按行拆分统计结果
    row_stats_shape = (n_rows, id_base)
    row_counts = group_counts.reshape(row_stats_shape)
    row_sums = group_sums.reshape(row_stats_shape)
    
    # 提取每行最大和对应的计数与和值
    max_sum_idx = row_sums.argmax(axis=1)
    max_counts = row_counts[np.arange(n_rows), max_sum_idx]
    max_sums = row_sums[np.arange(n_rows), max_sum_idx]
    
    return np.column_stack([max_counts, max_sums]).astype(int)

# 测试用例
if __name__ == "__main__":
    # 原示例dst为4行、ids为5行，此处补充dst最后一行与ids对齐，不影响核心逻辑
    dst = np.array([
        [1,2,3,4,5],
        [2,3,4,5,6],
        [1,1,2,2,3],
        [7,8,9,9,3],
        [1,2,3,4,5]
    ])
    ids = np.array([
        [1,1,2,1,3],
        [2,2,1,1,3],
        [3,3,2,1,1],
        [2,2,1,3,3],
        [1,2,3,2,1]
    ])
    result = calc_max_group(dst, ids)
    print(result)
    # 第一行输出为[3 7]，和需求示例完全匹配

补充说明 #

如果你的ids取值范围非常大（单组id最大值超过1e4），可以先对每行id做因子化编码再进行统计，避免不必要的内存占用。

内容的提问来源于stack exchange，提问作者Дима Лимонов