页面加载中…
AI023. 向量 Top-K 检索
中等通过率 0% · 提交 0 · 通过 0
华为 AI 岗Top-K堆点积
算法机考高频题型改编 · 聚焦「华为 AI 岗 / Top-K」 · 难度：中等
向量检索系统里有 n 个候选向量(编号从 1 开始)和一个查询向量,维度都是 d。候选向量与查询向量的相似度定义为两者的点积。请选出相似度最高的 k 个候选:相似度不同时,相似度大的优先;相似度相同时,编号小的优先。按这个优先顺序输出 k 个候选的编号。
这类题属于算法机考高频题型中「华为 AI 岗 / Top-K」方向的高频题型，通常考察对「华为 AI 岗 / Top-K」的建模能力与边界条件处理。掌握本题的解题思路后，可举一反三应对同类真题方向，稳步提升机考通过率。
输入描述
第一行输入三个整数 n d k。第二行输入 d 个整数,表示查询向量。接下来 n 行,每行输入 d 个整数,表示一个候选向量。
输出描述
输出一行 k 个整数,以单个空格分隔,表示选出的候选编号,按优先顺序排列。
示例
示例 1
输入示例
3 2 2 1 2 1 1 2 2 0 3
输出示例
2 3
说明：
相似度 6 并列,编号小的排前
示例 2
输入示例
2 3 2 0 0 0 5 5 5 -1 -1 -1
输出示例
1 2
说明：
查询向量为零向量时全部相似度为 0
时间限制 3000 ms · 内存限制 256 MB
本平台为独立第三方培训机构，与华为技术有限公司无任何关联；课程的服务内容与权益以购买协议为准，学习效果因个人情况而异。「华为 OD」「华为可信」等仅为对岗位与考试方向的客观描述，相关商标归各自权利人所有。
⚠️ 这道题的坑 · 易错点与反直觉
这些是真正决定能不能 AC、但通用题解里常被略过的点。
- 输出按(点积降序,编号升序)的优先序,不是编号升序
- 堆顶放的是当前 k 个里最差的,比较方向写反就留下了最差集合
解题思路
这道题在考什么
Top-K 检索是向量召回的骨架:点积打分 + 取前 k。课程第 2 天的堆和第 8 天的点积在这里合体,考点是并列规则和输出顺序这两条「可复现性纪律」,外加一个 64 位的数值意识。
优先序先写成一句话
「相似度大的优先;相似度相同时,编号小的优先。」写代码前把它翻译成排序键:(-点积, 编号) 升序。输出的是按这个优先序排列的 k 个编号——不是按编号升序,这条题面反复强调,示例 1 的输出 2 3 就是证据(相似度 6 并列,编号 2 在前)。
两条路:全排序或堆
- 全排序:算出 n 个 (点积, 编号),按键排序取前 k。O(n log n),n = 10 万完全可行,考场上这是最不容易写错的版本。
- 堆:维护大小为 k 的堆,只留当前最优的 k 个,O(n log k)。课程第 2 天 Top-K 的标准姿势,k 远小于 n 时更快。注意堆里比较的是「劣势方先出」:留大点积、同点积留小编号,堆顶放当前 k 个里最差的那个(点积最小、同点积编号最大),新候选比堆顶好才换。
时限给得宽裕,两条路都过;没把握就全排序,把正确性和并列规则钉死优先。
64 位点积
d ≤ 16、分量绝对值 ≤ 1000,单个点积最大 16×10⁶,int 装得下;但习惯上点积一律用 64 位累加——换个数据范围就不翻车,这个习惯值得养成。
零向量与全并列
查询是零向量时全部相似度为 0,退化成「取编号最小的 k 个」;全同相似度的用例专门在册,考的就是并列规则有没有落实。
易错点展开
- 输出按编号排序而不是优先序。
- 并列规则没进排序键,结果不可复现。
- 堆的比较方向写反:留下来的是最差的 k 个。
- 编号从 1 开始,读入时别悄悄变成 0 基。
参考实现怎么读
参考实现走堆版:维护大小 k 的小顶堆,堆中元素是 (点积, -编号) 的「优劣编码」——点积小的差、同点积编号大的差,堆顶恰好是当前 k 个里最差的。新候选比堆顶好就替换。结束后把堆里元素按 (-点积, 编号) 排序输出。如果这套编码绕,考场直接全排序:n log n 对 10 万完全无压力,少想一层反而快。
写完先自查这三组
1. 样例 1:点积 3、6、6,并列的 6 里编号 2 在前,输出 2 3——并列规则落实的直接证据。 2. k = n:全体输出,顺序 = (-点积, 编号) 全排序,验证输出序而不是编号序。 3. 全部候选与查询正交(点积全 0):退化为编号升序取前 k,并列规则独自决定一切。
数值与工程习惯
单个点积上限 16×10⁶,int 装得下,但参考实现仍用 64 位累加——检索场景里维度和量级说变就变,这个习惯在生产代码里同样保命。编号从 1 开始,读入循环里直接用 1 基枚举,避免输出时 +1 漏改。
和课程的连接
点积当相似度是第 8 天「距离与相似」的另一面,堆取 Top-K 是第 2 天的模板,第 16 天的检索语境把两者接上了生产场景——向量召回的核心循环和这道题一模一样,只是规模从 10 万变成上亿、单机变成分片。把这道题写顺,面试聊 RAG 召回时手上就有实感。
参考代码
# Top-K:排序键 (-点积, 编号),输出按优先序而非编号序
# 堆版维护 k 个当前最优,堆顶是最差者;点积用 64 位累加
import sys
import heapq
def solve() -> None:
data = sys.stdin.buffer.read().split()
if not data:
return
n = int(data[0]); d = int(data[1]); k = int(data[2])
at = 3
query = [int(data[at + j]) for j in range(d)]
at += d
# 键 = (-相似度, 编号):升序即优先顺序
keys = []
for i in range(1, n + 1):
s = 0
for j in range(d):
s += query[j] * int(data[at + j])
at += d
keys.append((-s, i))
top = heapq.nsmallest(k, keys)
sys.stdout.write(" ".join(str(idx) for _, idx in top) + "\n")
solve()
我的提交
登录后可查看你在本题的历史提交，以及每次的各用例通过情况。
