堆与 Top-K 问题

import sys
import heapq

data = sys.stdin.read().split()
n, d, k = int(data[0]), int(data[1]), int(data[2])
q = [int(x) for x in data[3:3 + d]]

heap = []                                  # 小顶堆：堆顶 = 当前 k 个候选里最先被淘汰的
pos = 3 + d
for idx in range(1, n + 1):                # 候选编号从 1 起
    v = data[pos:pos + d]
    pos += d
    score = sum(int(a) * b for a, b in zip(v, q))   # 点积；Python 整数不会溢出
    entry = (score, -idx)                  # 同分时编号大的更靠堆顶（先被淘汰），留下编号小的
    if len(heap) < k:
        heapq.heappush(heap, entry)
    elif entry > heap[0]:
        heapq.heapreplace(heap, entry)

# 堆里只是候选集合；按题目要求的优先顺序（相似度降序、编号升序）再排一次
kept = sorted(heap, key=lambda t: (-t[0], -t[1]))
print(" ".join(str(-neg) for _, neg in kept))

自测用例：题面两个示例、练习 1 的 k=n 输入、第 08 节错误表的 k=1 同分输入（含零查询向量、负分量）。

import sys
import heapq

data = sys.stdin.read().split()
n = int(data[0])
heap = [-int(x) for x in data[1:1 + n]]    # 取负入堆：小顶堆变成「最重的在堆顶」
heapq.heapify(heap)

while len(heap) >= 3:                      # 至少三块才开始一回合
    z = -heapq.heappop(heap)               # 弹出顺序：最重、次重、第三重
    y = -heapq.heappop(heap)
    x = -heapq.heappop(heap)               # 于是 x <= y <= z
    if x == y == z:
        rest = 0
    elif x == y:
        rest = z - y
    elif y == z:
        rest = y - x
    else:
        rest = abs((z - y) - (y - x))
    if rest > 0:
        heapq.heappush(heap, -rest)        # 非零才放回（记得取负）

print(-heap[0] if heap else 0)             # 剩两块输出较大者（堆顶）；剩一块输出它；空堆输出 0

自测用例：第 06 节两组手算例，以及 1 块、2 块、正好三块全熔（3 / 1 2 3）、4 / 5 5 5 1 四个边界（后两个在练习 5）。循环条件是堆大小 ≥ 3；rest > 0 才放回；输出前取负还原。

import heapq

def bottomk_ids(scores, k):
    heap = []
    for score, idx in scores:
        entry = (-score, -idx)           # 方向整体反过来：大顶堆；同分仍保编号小
        if len(heap) < k:
            heapq.heappush(heap, entry)
        elif entry > heap[0]:
            heapq.heapreplace(heap, entry)
    return sorted(-neg for _, neg in heap)                       # 本自测要求编号升序输出

assert bottomk_ids([(9, 1), (5, 2), (9, 3), (7, 4)], 2) == [2, 4]
assert bottomk_ids([(5, 1), (5, 3)], 1) == [1]

三处改动：① 得分取负（大顶堆，堆顶是当前 K 个里得分最高的）；② 同分规则跟着取负仍是 −编号；③ 输出规则按自测要求改为编号升序。替换条件 entry > heap[0] 不变。

import heapq

def topk_ids(scores, k):
    heap = []
    for score, idx in scores:
        entry = (score, -idx)            # 同分时编号大的更靠堆顶，先被淘汰
        if len(heap) < k:
            heapq.heappush(heap, entry)
        elif entry > heap[0]:
            heapq.heapreplace(heap, entry)
    return [-neg for _, neg in sorted(heap, key=lambda t: (-t[0], -t[1]))]   # 得分降序、同分编号升序

assert sorted(topk_ids([(9, 1), (5, 2), (9, 3), (7, 4)], 2)) == [1, 3]
assert topk_ids([(9, 1), (9, 3)], 1) == [1]
assert topk_ids([(3, 1), (6, 2), (6, 3)], 2) == [2, 3]        # AI023 题面示例 1 的得分

返回前按 (−得分, −(−编号)) 即得分降序、编号升序重排；第三条断言就是题面示例 1。