Python列表中查找使函数取最大/最小值的所有元素的最快方法

方案思路 #

仅需遍历列表1次，遍历过程中同步维护两个变量：

• 当前最大求值结果current_max
• 当前匹配最大求值结果的元素集合result

每个元素仅需调用1次求值函数，根据返回结果做对应处理：

• 若返回值大于current_max：更新current_max为当前返回值，清空result后加入当前元素
• 若返回值等于current_max：直接将当前元素追加到result中
• 若返回值小于current_max：跳过当前元素即可

代码实现 #

适配示例场景的代码 #

ll = [[1, 2], [1, 3], [2, 3], [1, 2, 3], [2, 3, 4]]

current_max = -1
res = []
for item in ll:
    item_len = len(item)
    if item_len > current_max:
        current_max = item_len
        res = [item]
    elif item_len == current_max:
        res.append(item)

print(res)

运行输出和原方案完全一致：
[[1, 2, 3], [2, 3, 4]]

通用封装版本（支持自定义求值函数） #

如果需要替换为其他高开销的求值函数，可以直接使用封装好的通用方法：

def get_all_max_elements(lst, eval_func):
    current_max = float('-inf')
    result = []
    for item in lst:
        eval_val = eval_func(item)
        if eval_val > current_max:
            current_max = eval_val
            result = [item]
        elif eval_val == current_max:
            result.append(item)
    return result

# 测试调用
ll = [[1, 2], [1, 3], [2, 3], [1, 2, 3], [2, 3, 4]]
print(get_all_max_elements(ll, len))

性能对比 #

• 原实现：需要2次全量遍历，每个元素调用2次求值函数，时间复杂度为O(2n*k)（n为列表长度，k为单次求值函数的开销）
• 优化后实现：仅需1次全量遍历，每个元素仅调用1次求值函数，时间复杂度为O(n*k)，列表规模越大、求值函数开销越高，性能提升越明显

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