如何从numpy相似度矩阵高效提取每个id前n个最相似元素降低内存占用

核心思路是跳过生成全量N²级别的Pandas结构，直接在Numpy阶段通过非全排序的高效取前值逻辑完成过滤，内存占用仅为原方案的1/n（n为保留的相似项数量，本例为10）。

方案1：已有sim矩阵的情况（仅需修改sim矩阵后的代码） #

import pandas as pd
import numpy as np
from scipy.spatial.distance import pdist, squareform

# 原有生成特征的代码保持不变
df = pd.read_csv('test.csv')
df = df[['id','value1', 'value2']]
id_list = df['id'].values  # 提前存好id映射，减少多余索引操作
feature_mat = df[['value1', 'value2']].values
sim = squareform(pdist(feature_mat, metric='cosine'))

n = 10  # 每个id保留的最相似项数量
# 自身和自身距离为0会排在第一位，所以多取1个后续过滤
k = n + 1

rows = []
for i in range(sim.shape[0]):
    # 用argpartition取前k个最小的索引，比全排序快3-10倍
    top_k_idx = np.argpartition(sim[i], k)[:k]
    # 仅对前k个做局部排序，保证结果顺序正确
    top_k_idx = top_k_idx[np.argsort(sim[i][top_k_idx])]
    # 跳过自身匹配项，取后面n个有效结果
    for j in top_k_idx[1:]:
        rows.append([id_list[i], id_list[j], sim[i][j]])

# 直接生成最终结果表，无全量N²中间数据
table = pd.DataFrame(rows, columns=['id_1', 'id_2', 'distance'])

方案2：超大数据集优化（不生成全量sim矩阵，内存占用再降N倍） #

如果数据量极大，连N*N的sim矩阵都无法存入内存，可以逐行计算相似度，算完即处理不用存储全矩阵：

import pandas as pd
import numpy as np
from sklearn.metrics.pairwise import cosine_distances

df = pd.read_csv('test.csv')
id_list = df['id'].values
feature_mat = df[['value1', 'value2']].values
n = 10
k = n + 1

rows = []
for i in range(feature_mat.shape[0]):
    # 单次计算当前行和所有样本的相似度，仅占用1*N的内存
    cur_dist = cosine_distances(feature_mat[i:i+1], feature_mat)[0]
    top_k_idx = np.argpartition(cur_dist, k)[:k]
    top_k_idx = top_k_idx[np.argsort(cur_dist[top_k_idx])]
    for j in top_k_idx[1:]:
        rows.append([id_list[i], id_list[j], cur_dist[j]])

table = pd.DataFrame(rows, columns=['id_1', 'id_2', 'distance'])

关键优化点说明 #

• 用np.argpartition替代全排序，仅提取需要的前n个最小距离，时间复杂度从O(N²logN)降至O(N²)，大数据量下提速明显
• 全程用列表缓存结果，最后一次性转DataFrame，避免生成N²级别的中间Pandas结构，内存占用从O(N²)降至O(N*n)，n=10时内存仅为原方案的千分之一甚至更低（N越大优化效果越明显）
• 方案2完全放弃存储全量相似度矩阵，内存占用进一步降至O(N)，适配百万级以上样本的相似度计算场景

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