快速遍历多组元组列表查找每组内float最大值对应元组的高效实现方法

输入与目标 #

• 输入为数万个长度不一的子元组列表，每个子列表的元素均为(int, float)格式的元组对，示例（已修正原示例语法错误）：

[
    [(0, 0.3792), (3, 0.5796)],
    [(0, 0.9365), (1, 0.0512), (18, 0.0123)],
    [(13, 0.8642)],
    [(0, 0.6249), (1, 0.01), (2, 0.01), (3, 0.01), (4, 0.01), (5, 0.01)]
]

• 处理要求：对每个子元组列表，返回其中float值最大的完整(int, float)元组对，例如上述示例中第一个列表输出(3, 0.5796)，第四个列表输出(0, 0.6249)。

现有实现与遇到的问题 #

已落地的两种实现方案 #

1. 单列表转numpy结构化数组计算

def get_max(doc: List[Tuple[int, float]]) -> Tuple[int, float]:
    topic_prob_array = np.array(doc, dtype=np.dtype('int,float'))
    return topic_prob_array['f0'][np.argmax(topic_prob_array['f1'])], np.max(topic_prob_array['f1'])

2. 元组拆包后基于numpy计算

def get_max(doc: List[Tuple[int, float]]) -> Tuple[int, float]:
    ids, probabilities = zip(*doc)
    return ids[np.argmax(probabilities)], np.max(probabilities)

现存问题 #

• 尝试对单列表处理逻辑做numpy矢量化封装时，先后用np.vectorize、np.fromfunc实现，但对输入输出匹配格式判断错误，始终无法正常运行矢量化版本。
• 上述两种非矢量化方案性能基本持平，处理约8万组元组列表时耗时约1分10秒，希望尽可能压缩处理时长。

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