向量 Top-K 检索（AI023）- 算法机考高频题型 · 华为 AI 岗 / Top-K题解 | AlgoMooc

AI023. 向量 Top-K 检索

中等

通过率 81% · 提交 16 · 通过 13

华为 AI 岗Top-K堆点积

这道 OJ 的配套动画

向量 Top-K 检索 · 对应本题的思路、过程和代码

算法机考高频题型改编 · 聚焦「华为 AI 岗 / Top-K」 · 难度：中等

向量检索系统里有 n 个候选向量(编号从 1 开始)和一个查询向量,维度都是 d。候选向量与查询向量的相似度定义为两者的点积。请选出相似度最高的 k 个候选:相似度不同时,相似度大的优先;相似度相同时,编号小的优先。按这个优先顺序输出 k 个候选的编号。

这类题属于算法机考高频题型中「华为 AI 岗 / Top-K」方向的高频题型，通常考察对「华为 AI 岗 / Top-K」的建模能力与边界条件处理。掌握本题的解题思路后，可举一反三应对同类真题方向，稳步提升机考通过率。

输入描述

第一行输入三个整数 n d k。第二行输入 d 个整数,表示查询向量。接下来 n 行,每行输入 d 个整数,表示一个候选向量。

输出描述

输出一行 k 个整数,以单个空格分隔,表示选出的候选编号,按优先顺序排列。

数据范围

1 <= n <= 100000;1 <= d <= 16;1 <= k <= n;所有分量的绝对值不超过 1000。

示例

示例 1

输入示例

输出示例

2 3

说明：

相似度 6 并列,编号小的排前

示例 2

输入示例

输出示例

1 2

本平台为独立第三方培训机构，与华为技术有限公司无任何关联；课程的服务内容与权益以购买协议为准，学习效果因个人情况而异。「华为 OD」「华为可信」等仅为对岗位与考试方向的客观描述，相关商标归各自权利人所有。

问小欧 · 图解算法年卡¥99 开通 →

题库上一题下一题

未登录也能直接「运行」调试;「提交」判题需微信登录,你写的代码会自动保留、登录后回到本题继续。

字号13Ctrl+/ 注释 · Ctrl+↵ 提交

1
2

# 在此编写你的代码

写完代码点「提交」，将对全部测试用例判题。

⚠️ 这道题的坑 · 易错点与反直觉

这些是真正决定能不能 AC、但通用题解里常被略过的点。

输出按(点积降序,编号升序)的优先序,不是编号升序
堆顶放的是当前 k 个里最差的,比较方向写反就留下了最差集合

解题思路

先计算查询向量与每个候选向量的点积。候选的完整优先级是“点积更大优先，点积相同时编号更小优先”，可表示为排序键 (-score, id)。直接排序全部候选即可通过；也可以维护大小为 k 的堆，使堆顶始终是当前结果中优先级最低的候选，读入新候选时只在它更优的情况下替换堆顶。最后仍需按完整优先级排序堆中元素后输出编号。

正确性说明

固定大小堆始终保存已处理候选中的前 k 名：不足 k 个时直接加入；达到 k 个后，新候选只有优于当前第 k 名时才会替换它。归纳处理全部候选后，堆中正好是全局前 k 名。最终排序落实题目要求的输出顺序。

复杂度分析

点积计算需要 O(nd) 时间，堆维护需要 O(n log k) 时间，最终排序需要 O(k log k)；空间为 O(k)。点积使用 64 位整数累加。

边界与常见错误

查询向量为零向量时，所有点积相同，应输出编号最小的 k 个候选。不要把最终结果改按编号单独排序，也不要忽略点积相同情况下的编号规则。

参考代码

# 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()