有没有办法进一步优化 Python 的 heapq.nlargest 以选择前 N 个项目？

Posted 2023-03-06

【中文标题】有没有办法进一步优化 Python 的 heapq.nlargest 以选择前 N 个项目？

【英文标题】：Is there any way to further optimize Python's heapq.nlargest for selecting top N items?

【发布时间】：2016-01-24 09:04:03

【问题描述】：

我使用heapq.nlargest 选择前 N 个项目，它占用了 98% 的运行时间（参见第 51 行）：

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
40                                           @profile
41                                           def gen_submit(index_to_pri, index_to_sec, exclude_set, pri_mat, sec_mat, gen_count):
42         1           33     33.0      0.0      print('gen_submit')
43         1           87     87.0      0.0      f = open('../submission.txt', 'w')
44        16           28      1.8      0.0      for i, pri in enumerate(index_to_pri):
45        16          369     23.1      0.0          print('generate recommendation for %d-th primary object' % i)
46        16          103      6.4      0.0          recommend_sec = []
47        16           25      1.6      0.0          exclude = exclude_set[pri]
48        16        68215   4263.4      1.3          rating_vector = numpy.dot(pri_mat[i], sec_mat.T)
49                                                   # extract top N
50        16          102      6.4      0.0          N = 500 + len(exclude_set[pri])
51        16      4988735 311795.9     98.2          top_N_indexed_rating = heapq.nlargest(N, enumerate(rating_vector), key = lambda x: x[1]))
52        15          181     12.1      0.0          top_N_j = map(lambda x: x[0], top_N_indexed_rating)
53      7501         6229      0.8      0.1          for j in top_N_j:
54      7501         4812      0.6      0.1              if not index_to_sec[j] in exclude:
55      7500         6135      0.8      0.1                  recommend_sec.append(str(j))
56      7500         4943      0.7      0.1                  if len(recommend_sec) >= 500: break
57        15          293     19.5      0.0          f.write(' '.join(recommend_sec) + '\n')
58                                               f.close()

如何进一步优化这个单一的操作？

【问题讨论】：

是否可以逐步构建 N 个最高值的列表？似乎最大只对较小的 n 值有效。

@PaulRooney 你的意思是遍历序列 500 次（在我的情况下）？另外我需要维护一个集合来排除我已经找到的前k个项目，这看起来很乏味......更糟糕的是复杂度从O（nlogk）变成了O（nk）。

FWIW，文件使用with；不要手动close他们。

您需要它们按顺序排列吗？您能提供示例数据吗？

exclude 和 rating_vector 通常有多大？

【参考方案1】：

新答案

如果你在top_N_j内不需要下单，试试

top_N_j = rating_vector.argpartition(len(rating_vector) - N)[-N:]

否则之后用

排序

top_N_j = top_N_j[numpy.argsort(rating_vector[top_N_j])]

我认为这比你给的时间少了大约 30 到 50 倍。

---

旧答案

我想这太明显了，我很可能完全忽略了这一点，但是

heapq.nlargest(N, enumerate(...))

将只取最后一个N 元素，由它们的索引标记，以相反的顺序。然后，您仅将其用于

top_N_j = map(lambda x: x[0], top_N_indexed_rating)

这会变成单独的索引。

看来你想要的是

end = len(...)
start = max(0, end - N)
top_N_j = reversed(range(start, end))

（尽管我必须承认自己对你所做的事情感到非常困惑。）

【讨论】：

感谢您指出我的错误。实际上我在我的电脑上更正了它，但没有在这里更新我的代码。现在更新了。

我所做的就是获取top N元素的索引，一句话。

拍得好！此操作的时间成本已从 97% 降低到 27%，收缩系数为 7.7，但比您想象的要低得多。

看起来分区的复杂度低于“堆过滤器”，后者的成本为 O(nlogk)。