基于相似度分数生成排名型DataFrame的高效实现问询

现有一个方阵型DataFrame，索引和列名均为图片文件名，存储了各图片间的相似度分数。需要为每张图片提取相似度排名Top50的图片，生成一个300000×50的紧凑排名型DataFrame。

原方案通过循环遍历每一列排序取前50，但处理30万条数据时效率极低，现寻求更高效的实现方法。

原始DataFrame示例 #

name_A.jpg     name_B.jpg     name_C.jpg     name_D.jpg   ...
name_A.jpg     1.000000       0.725098       0.291748       0.444336
name_B.jpg     0.725098       1.000000       0.255371       0.482178
name_C.jpg     0.291748       0.255371       1.000000       0.382812
name_D.jpg     0.444336       0.482178       0.382812       1.000000
...

期望输出示例 #

0              1              2              3              ....          49
name_A.jpg     name_B.jpg     name_D.jpg     name_N.jpg     name_O.jpg     
name_B.jpg     name_A.jpg     name_D.jpg     ...            ...
name_C.jpg     ...            ...            ...            ...
name_D.jpg     name_B.jpg     name_A.jpg     ...            ...
...

方法1：NumPy向量化TopK筛选（极致性能） #

针对30万行的超大数据集，用NumPy的argpartition实现O(n)级别的TopK选择，比全排序效率高出一个数量级，完全避免循环开销。

import numpy as np
import pandas as pd

# 假设原始相似度矩阵为sim_df
top_n = 50
# 获取所有图片文件名（列名）
image_names = sim_df.columns.values

# 对每行取Top50+1个最大相似度的索引（+1是为了排除自身）
# 用负号是因为argpartition默认取最小的，取负后等价于取最大的
top_indices = np.argpartition(-sim_df.values, top_n + 1, axis=1)[:, 1:top_n+1]

# 将索引映射为对应的图片文件名
top_image_names = image_names[top_indices]

# 生成结果DataFrame
result_df = pd.DataFrame(top_image_names, index=sim_df.index, columns=range(top_n))

方法2：Pandas内置方法（简洁易读） #

如果对性能要求不是极致，用Pandas的nlargest结合行级apply实现，代码更直观：

top_n = 50

def extract_top50(row):
    # 取Top50+1后跳过第一个（自身），返回剩余的Top50文件名
    return row.nlargest(top_n + 1).index[1:]

# 应用函数并转换为DataFrame
result_df = sim_df.apply(extract_top50, axis=1).apply(pd.Series)

性能对比 #

• 方法1的NumPy向量化操作是最优选择，处理30万行数据的速度比循环列快100倍以上。
• 方法2的apply本质是行级循环，但比手动遍历列高效，适合快速实现或中小规模数据。

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