Python如何从字典的多个等长列表中按索引获取最大值对应键

给定结构如下的字典，字典中所有list_前缀的列表长度完全相同：

my_dict = {'list_1': [1,2,3,4,5],
           'list_2': [3,4,1,0,0],
           'list_3': [1,2,3,0,6]}

需要生成与列表等长的结果列表final_list，规则为：

• 每个索引位置的元素，为该位置上取值最大的列表对应的字典键
• 若同一位置存在多个并列最大值，将所有符合条件的键用|拼接后作为该位置的元素
• 对应上述示例的期望输出为：['list_2', 'list_2', 'list_1|list_3', 'list_1', 'list_3']

以下两种实现均达到该问题的理论最优时间复杂度O(n*m)（n为单列表长度，m为字典内列表数量），区别仅在于适用场景不同。

方案1：原生Python实现（无依赖，通用场景首选） #

直接利用Python内置的zip函数按索引聚合所有列表的同位置元素，逐位匹配最大值对应的键，内置函数均为C层实现，执行效率高，无额外依赖：

def build_final_list(input_dict):
    # 若字典存在非list_前缀的无关键，可打开下面这行过滤
    # input_dict = {k:v for k,v in input_dict.items() if k.startswith('list_')}
    keys = list(input_dict.keys())
    final_list = []
    for pos_values in zip(*input_dict.values()):
        current_max = max(pos_values)
        match_keys = [k for k, v in zip(keys, pos_values) if v == current_max]
        final_list.append('|'.join(match_keys))
    return final_list

# 测试运行
my_dict = {'list_1': [1,2,3,4,5],
           'list_2': [3,4,1,0,0],
           'list_3': [1,2,3,0,6]}
print(build_final_list(my_dict))
# 输出：['list_2', 'list_2', 'list_1|list_3', 'list_1', 'list_3']

方案2：NumPy向量化实现（超大数据量场景提速） #

当单列表长度达到十万级以上、字典内列表数量较多时，用NumPy的向量化运算替代Python层显式循环，性能可提升数倍到数十倍：

import numpy as np

def build_final_list_np(input_dict):
    # 若字典存在非list_前缀的无关键，可打开下面这行过滤
    # input_dict = {k:v for k,v in input_dict.items() if k.startswith('list_')}
    keys = list(input_dict.keys())
    value_arr = np.array(list(input_dict.values()))
    col_max = value_arr.max(axis=0)
    final_list = []
    for col in range(value_arr.shape[1]):
        match_idx = np.where(value_arr[:, col] == col_max[col])[0]
        final_list.append('|'.join(keys[i] for i in match_idx))
    return final_list

注意：该方案需要安装NumPy依赖，小数据量下因为数组初始化的额外开销，运行速度反而低于原生Python方案，仅建议大数据量场景使用。

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