KV 缓存与束搜索

import math

def rope(x, p, base=10000):
    d = len(x)
    y = [0.0] * d
    for i in range(d // 2):                      # 维度两两配对 (2i, 2i+1)
        theta = base ** (-2 * i / d)             # 第 i 对的角频率
        a = p * theta                            # 位置 p 的旋转角
        c, s = math.cos(a), math.sin(a)
        y[2 * i] = x[2 * i] * c - x[2 * i + 1] * s
        y[2 * i + 1] = x[2 * i] * s + x[2 * i + 1] * c
    return y

# AI039 样例 1：d=4、p=3
y = rope([1.00, 0.50, -2.00, 0.25], 3)
assert [round(v, 4) for v in y] == [-1.0606, -0.3539, -2.0066, 0.1899]
# p=0：旋转角全 0，原样输出
assert [round(v, 4) for v in rope([1.00, 0.50, -2.00, 0.25], 0)] == [1.0, 0.5, -2.0, 0.25]
# 旋转不改变每一对的长度：|y[2i], y[2i+1]| == |x[2i], x[2i+1]|
assert abs(math.hypot(y[0], y[1]) - math.hypot(1.0, 0.5)) < 1e-9

断言覆盖 AI039 样例 1、p=0 原样输出，以及「旋转不改变每一对的长度」。θ 用浮点幂 base ** (-2 * i / d)。

import sys


def common_prefix(a, b):
    k = 0
    while k < len(a) and k < len(b) and a[k] == b[k]:
        k += 1
    return k


def main():
    data = list(map(int, sys.stdin.read().split()))
    if not data:
        return
    s, q = data[0], data[1]
    pos = 2
    loads, seqs = [], []
    for _ in range(s):
        load, length = data[pos], data[pos + 1]
        seqs.append(data[pos + 2:pos + 2 + length])
        loads.append(load)
        pos += 2 + length
    out = []
    for _ in range(q):
        length = data[pos]
        req = data[pos + 1:pos + 1 + length]
        pos += 1 + length
        best = -1
        best_key = None
        for i in range(s):
            k = common_prefix(seqs[i], req)
            key = (-k, loads[i], i)              # 前缀长者优先 → 负载小者优先 → 编号小者优先
            if best_key is None or key < best_key:
                best, best_key = i, key
        matched = -best_key[0]
        out.append(f"{best + 1} {matched}")      # 槽位编号从 1 起
        loads[best] += length - matched          # 负载加未命中的 token 数
        seqs[best] = req                         # 缓存序列换成本次请求
    print("\n".join(out))


main()

每个请求逐槽位算前缀 O(S·len)，S=Q=500、len=60 时最多 500 × 500 × 60 = 1.5×10⁷ 次词元比较。用题面示例核对：输出 2 2 / 3 1；提交前另在目标语言中检查最大规模耗时。

import bisect
import sys


def main():
    data = list(map(int, sys.stdin.read().split()))
    if not data:
        return
    n, q = data[0], data[1]
    base, fixed, kv, attn = data[2:6]
    lengths = sorted(data[6:6 + n])              # 先短后长：同样的 c，最短的 c 个显存最小
    pos = 6 + n
    need = [base]                                # need[c] = 放最短 c 个请求所需显存
    total = 0
    for c in range(1, n + 1):
        total += lengths[c - 1]
        need.append(base + fixed * c + kv * c * lengths[c - 1] + attn * total)
    out = []
    for _ in range(q):
        budget = data[pos]
        pos += 1
        c = bisect.bisect_right(need, budget) - 1   # need 单调不减，二分找最大的 c
        out.append(str(max(0, c)))
    print("\n".join(out))


main()

排序和 need 数组各计算一次，每个预算只做一次二分。用题面示例核对：输出 1 / 2 / 3。

import heapq
import sys


def main():
    data = sys.stdin.read().split()
    if not data:
        return
    C, q = int(data[0]), int(data[1])
    cache = {}                                   # id -> [score, 加入序号]
    heap = []                                    # (score, 加入序号, id)，惰性删除
    order = 0
    for k in range(q):
        tid, score = int(data[2 + 2 * k]), int(data[3 + 2 * k])
        if tid in cache:                         # 已在缓存：只改得分，加入顺序不变
            cache[tid][0] = score
            heapq.heappush(heap, (score, cache[tid][1], tid))
            continue
        if len(cache) == C:                      # 已满：淘汰得分最低、并列最早加入者
            while True:
                sc, ar, vid = heapq.heappop(heap)
                ent = cache.get(vid)
                if ent is not None and ent[0] == sc and ent[1] == ar:   # 跳过过期堆项
                    del cache[vid]
                    break
        cache[tid] = [score, order]
        heapq.heappush(heap, (score, order, tid))
        order += 1
    if not cache:
        print("empty")
        return
    rows = sorted(cache.items(), key=lambda kv: kv[1][1])
    print("\n".join(f"{tid} {ent[0]}" for tid, ent in rows))


main()

堆里的过期项在弹出时丢弃，每条记录至多入堆一次、每个堆项至多弹出一次，总量 O(q log q)，不再逐次扫描整个缓存。用第 07 节的八条记录核对：输出 2 10 / 4 8。

import sys


def main():
    data = list(map(int, sys.stdin.read().split()))
    if not data:
        return
    T, V, B = data[0], data[1], data[2]
    pos = 3
    first = data[pos:pos + V]
    pos += V
    mats = []
    for _ in range(T - 1):
        mats.append([data[pos + i * V:pos + (i + 1) * V] for i in range(V)])
        pos += V * V
    beam = [(0, ())]                             # (总分, 序列)
    for t in range(T):
        cands = []
        for score, seq in beam:
            row = first if t == 0 else mats[t - 1][seq[-1]]
            for j in range(V):
                cands.append((score + row[j], seq + (j,)))
        cands.sort(key=lambda c: (-c[0], c[1]))  # 总分降序，同分序列字典序升序
        beam = cands[:B]
    score, seq = beam[0]
    print(score)
    print(" ".join(map(str, seq)))


main()

每步最多 B×V = 400 个候选，排序后截断。用题面示例核对：输出 6 / 1 1。

import math

def rope(x, p, base=10000):
    d = len(x)
    y = [0.0] * d
    for i in range(d // 2):                      # 维度两两配对 (2i, 2i+1)
        theta = base ** (-2 * i / d)             # 第 i 对的角频率
        a = p * theta                            # 位置 p 的旋转角
        c, s = math.cos(a), math.sin(a)
        y[2 * i] = x[2 * i] * c - x[2 * i + 1] * s
        y[2 * i + 1] = x[2 * i] * s + x[2 * i + 1] * c
    return y

# AI039 样例 1：d=4、p=3
y = rope([1.00, 0.50, -2.00, 0.25], 3)
assert [round(v, 4) for v in y] == [-1.0606, -0.3539, -2.0066, 0.1899]
# p=0：旋转角全 0，原样输出
assert [round(v, 4) for v in rope([1.00, 0.50, -2.00, 0.25], 0)] == [1.0, 0.5, -2.0, 0.25]
# 旋转不改变每一对的长度：|y[2i], y[2i+1]| == |x[2i], x[2i+1]|
assert abs(math.hypot(y[0], y[1]) - math.hypot(1.0, 0.5)) < 1e-9

见第 09 节展开区（同一份代码，最后一条断言就是本题）。