P5467. 全局位与窗口注意力

场次：2026年9月18日 · 来源：CodeFun2000 P5467

题目描述

给定三个整数矩阵 $Q$、$K$、$V$，尺寸均为 $L\times D$，以及 $G$ 个全局位置，它们构成集合 $P$。

对位置 $i$，先确定可参与注意力的下标集合 $C_i$：

$$C_i= \begin{cases} \{0,1,\ldots,L-1\}, & i\in P\\ P\cup\{j\mid 0\le j<L,\ |i-j|\le W\}, & i\notin P \end{cases}$$

集合内同一位置只计一次。

对 $j\in C_i$，注意力分数定义为

$$score(i,j)=\sum_{x=0}^{D-1}Q[i][x]\cdot K[j][x]$$

只保留满足 $score(i,j)>0$ 的位置，再从中选出至多 $T$ 个分数最大的候选，得到 $A_i$。分数相同时，下标 $j$ 较小者优先。若正分数候选个数不超过 $T$，则全部保留。

输出矩阵 $O$：当 $A_i$ 非空时

$$O[i][x]=\sum_{j\in A_i}score(i,j)\cdot V[j][x]\qquad(0\le x<D)$$

当 $A_i$ 为空时，$O[i]$ 的每一维都是 $0$。

请输出最终的 $L\times D$ 矩阵 $O$。

输入格式

第一行五个整数：

L D G W T

依次为序列长度、特征维数、全局位置个数、局部窗口半径、每个位置最多保留的候选数。

若 $G>0$，下一行 $G$ 个互不相同的整数：

p1 p2 ... pG

表示全局位置，下标从 $0$ 开始。

若 $G=0$，则没有这一行，接下来直接读矩阵 $Q$。

随后 $L$ 行，每行 $D$ 个整数，为矩阵 $Q$。

随后 $L$ 行，每行 $D$ 个整数，为矩阵 $K$。

随后 $L$ 行，每行 $D$ 个整数，为矩阵 $V$。

输出格式

输出 $L$ 行，每行 $D$ 个整数，表示矩阵 $O$。

同一行相邻整数之间用一个空格分隔。

示例

示例 1

输入：

5 2 2 1 2
0 3
1 0
1 0
2 0
0 1
1 1
2 0
1 0
3 0
0 2
-1 1
1 1
2 2
3 3
4 4
5 5

输出：

11 11
11 11
22 22
13 13
10 10

说明：

全局位置集合为 $P=\{0,3\}$。

对 $i=1$（非全局）：$C_1=\{0,1,2,3\}$，分数为 $2,1,3,0$。丢掉非正分数后，按 $T=2$ 取位置 $2$ 与 $0$：

$$O[1]=3\times[3,3]+2\times[1,1]=[11,11]$$

对全局位置 $i=3$：$C_3=\{0,1,2,3,4\}$，正分数只有 $2$（对应 $j=3$）和 $1$（对应 $j=4$）：

$$O[3]=2\times[4,4]+1\times[5,5]=[13,13]$$

对 $i=4$：$C_4=\{0,3,4\}$，分数为 $2,2,0$。$0$ 被丢掉，两个分数 $2$ 并列时先取较小下标 $0$，再取 $3$：

$$O[4]=2\times[1,1]+2\times[4,4]=[10,10]$$

示例 2

输入：

3 1 0 0 1
1
2
-1
-2
3
1
10
20
30

输出：

0
120
0

说明：

本例 $G=0$，没有全局位置行。$W=0$，每个位置的 $C_i$ 只含自己。

$i=0$ 的分数为 $-2$，不进入 $A_0$，输出 $0$。

$i=1$ 的分数为 $6$，因此

$$O[1]=6\times 20=120$$

$i=2$ 的分数为 $-1$，输出 $0$。

提示

$1\le L\le 10000$

$1\le D\le 64$

$0\le G\le \min(L,32)$

$0\le W\le 128$

$1\le T\le \min(L,16)$

$0\le p_i<L$

全局位置两两不同

$-10\le Q[i][j],K[i][j],V[i][j]\le 10$

全局位置的输入顺序不影响结果。

矩阵与序列下标均从 $0$ 开始。

中间结果和最终结果均在 $64$ 位有符号整数范围内。

---

解题思路

第一步：理解问题本质

注意力机制的通用骨架是：每个位置 $i$ 拿自己的 query 向量去和若干位置的 key 向量打分，
再用这些分数对相应位置的 value 向量加权求和。本题在这个骨架上做了四处改动：

1. 感受野不一样。普通位置的候选只有「左右各 $W$ 的窗口」加上「少数几个全局位置」；
   而全局位置本身就是负责汇总远距离信息的，所以它能看见整条序列。
   这正是 Longformer、BigBird 这类稀疏注意力的做法：绝大多数 token 只看邻居，
   少数 global token 看全局，从而把 $O(L^2)$ 的注意力降到近似线性。
2. 不做 softmax，也不除以 $\sqrt{D}$。点积算出来是多少，就直接当权重用。
3. 只保留正分数，并且最多取 $T$ 个。分数为 $0$ 或负数的位置一律不参与，
   相当于一种「硬阈值 + 截断」的稀疏化。
4. 并列时按下标打破平局：分数相同，$j$ 小的优先。

输出与输入同维：每个位置得到一个 $D$ 维向量 $O[i]$；一个候选都没有时整行为 $0$。

第二步：暴力解法

最直接的做法是对每个 $i$ 都把全部 $L$ 个位置当成候选：

for i in range(L):
    scored = []
    for j in range(L):                         # 不区分全局位置与普通位置
        s = 0
        for x in range(D):                     # 点积
            s += Q[i][x] * K[j][x]
        if s > 0:
            scored.append((s, j))

    # 分数降序、下标升序
    scored.sort(key=lambda pair: (-pair[0], pair[1]))
    A = scored[:T]

    row = [0] * D
    for s, j in A:
        for x in range(D):
            row[x] += s * V[j][x]
    O[i] = row

这个写法是对的，但每个位置都要和全部 $L$ 个 key 打分，总计算量是
$O(L^2D)$。$L=10^4$、$D=64$ 时约 $6.4\times10^9$ 次乘加，
而且如果先把完整的 $L\times L$ 分数矩阵存下来，光内存就是 $10^8$ 个数。

第三步：优化解法——只算真正需要的候选

观察 $C_i$ 的定义，绝大多数位置根本不需要看全序列：

• 全局位置只有 $G$ 个（$G\le 32$），它们才需要扫完整条序列，代价是 $G\cdot L$ 个点积；
• 其余 $L-G$ 个普通位置，候选是「窗口」并上「全局位置」，
  数量不超过 $(2W+1)+G$。

于是总分数的量级从 $L^2$ 降到

$$(L-G)\cdot(2W+1+G)\;+\;G\cdot L .$$

以最坏情况 $L=10^4, W=128, G=32$ 估算，普通位置每行约 $257+32\approx289$ 个候选，
总数约 $2.9\times10^6$ 个点积，大约只有完整注意力（$L^2=10^8$ 个）的三十分之一。

还有一个容易忽略的点：分数不能跨行复用。滑动窗口注意力里可以靠增量更新复用
窗口内的分数，但这里每换一行，query 就整个换掉了，点积的结果全都不同，
所以只能逐行重算。

第四步：最优解法——完整实现细节

(1) 窗口要截断，不要补零。 位置 $i$ 的窗口用左闭右开区间表示：

lo = max(0, i - W)
hi = min(L, i + W + 1)      # 右端开，所以是 i + W + 1

靠近两端的行，窗口天然变短，不需要任何特殊处理，循环 range(lo, hi) 即可。

(2) 候选集合要去重。 非全局行的候选是「窗口 ∪ 全局位置」。
如果某个全局位置恰好落在窗口里，它会被两条路径各算一次。
代码里先遍历窗口，再补上窗口外的全局位置：

for j in range(lo, hi):
    ...                       # 窗口内的位置
for j in global_pos:
    if j < lo or j >= hi:     # 只补窗口外的，避免重复
        ...

重复计入的后果不是「多算一个候选」这么简单：分数会被重复加进 $O[i]$，结果直接错。

(3) 只留正分数。 题目要求 $score(i,j)>0$，所以 $0$ 也要丢掉。
这一步同时起到了剪枝作用：全为负分的行直接输出零行。

(4) 取前 $T$ 个，并列看下标。 用 (分数, 下标) 做候选，
排序键写成 (-score, j)，升序排序后自然满足「分数从大到小、同分下标从小到大」。
如果正分候选不足 $T$ 个，就全部保留，不需要补齐。

(5) 加权求和得到输出行。 对每个入选的 $j$，把 $score(i,j)$ 乘到 $V[j]$ 的每一维上：

row = [0] * D
for score, j in candidates:
    vj = V[j]
    for x in range(D):
        row[x] += score * vj[x]

(6) 输入解析。 $G=0$ 时没有全局位置那一行，输入的行数会变。
按「token 流」一次性读入再顺序消费，天然不受行结构影响：

tokens = sys.stdin.buffer.read().split()

(7) 点积的实现。 点积是整道题唯一的热点。Python 3.12 起标准库提供了
C 实现的 math.sumprod，比 sum(map(mul, a, b)) 快将近一倍；
旧版本 Python 用后者，两者结果完全等价。代码里用一个 dot 名字统一起来，
主体逻辑不需要关心用的是哪一份实现。

---

完整代码实现

import sys

try:
    # Python 3.12 起标准库提供了 C 实现的点积，比手写循环快将近一倍
    from math import sumprod as dot
except ImportError:                     # 旧版本 Python 的等价写法
    from operator import mul

    def dot(a, b):
        """返回两个等长向量的点积。"""
        return sum(map(mul, a, b))


def solve():
    tokens = sys.stdin.buffer.read().split()
    if not tokens:
        return

    pos = 0
    L = int(tokens[pos])
    D = int(tokens[pos + 1])
    G = int(tokens[pos + 2])
    W = int(tokens[pos + 3])
    T = int(tokens[pos + 4])
    pos += 5

    # is_global 用布尔数组存，判断「i 是不是全局位置」是 O(1)；
    # global_pos 保存原始下标列表，非全局行要把它们补进候选集合。
    is_global = bytearray(L)
    global_pos = []
    for _ in range(G):
        p = int(tokens[pos])
        pos += 1
        is_global[p] = 1
        global_pos.append(p)

    def read_matrix(rows):
        """从 token 流中读出 rows 行、每行 D 个整数的矩阵。"""
        nonlocal pos
        mat = []
        for _ in range(rows):
            mat.append([int(x) for x in tokens[pos:pos + D]])
            pos += D
        return mat

    Q = read_matrix(L)
    K = read_matrix(L)
    V = read_matrix(L)
    # 三个矩阵已经转换成整数，token 列表不再需要，提前释放可以省下大量内存。
    del tokens

    output_lines = []
    for i in range(L):
        qi = Q[i]
        candidates = []          # 只放正分候选，元素为 (分数, 下标)

        if is_global[i]:
            # 全局位置的感受野是整条序列
            for j in range(L):
                score = dot(qi, K[j])
                if score > 0:
                    candidates.append((score, j))
        else:
            # 窗口在序列两端会被截断，而不是补零
            lo = i - W
            if lo < 0:
                lo = 0
            hi = i + W + 1
            if hi > L:
                hi = L

            for j in range(lo, hi):
                score = dot(qi, K[j])
                if score > 0:
                    candidates.append((score, j))

            # 全局位置可能落在窗口内，那它已经在上面被算过一次，
            # 这里只补窗口外的那些，避免同一位置被重复计入。
            for j in global_pos:
                if j < lo or j >= hi:
                    score = dot(qi, K[j])
                    if score > 0:
                        candidates.append((score, j))

        if not candidates:
            # 没有正分候选时整行为 0，这是题目明确规定的
            output_lines.append(" ".join(["0"] * D))
            continue

        if len(candidates) > T:
            # 分数降序、分数相同时下标升序：升序排序键 (-score, j) 正好表达这个顺序
            candidates.sort(key=lambda pair: (-pair[0], pair[1]))
            candidates = candidates[:T]

        row = [0] * D
        for score, j in candidates:
            vj = V[j]
            # O[i][x] += score(i,j) * V[j][x]
            for x in range(D):
                row[x] += score * vj[x]

        output_lines.append(" ".join(map(str, row)))

    sys.stdout.write("\n".join(output_lines))


if __name__ == "__main__":
    solve()

NumPy 版本

以下版本使用 NumPy。只有在评测环境允许第三方库时才能提交；CodeFun2000 的普通 Python 环境不一定安装 NumPy。

import sys

import numpy as np


def solve():
    """带全局位置的窗口注意力：候选筛选照旧，点积与加权求和全部交给 NumPy。"""
    tokens = sys.stdin.buffer.read().split()
    if not tokens:
        return

    L, D, G, W, T = (int(t) for t in tokens[:5])
    cursor = 5

    # G = 0 时输入里没有全局位置那一行；按 token 流顺序消费，天然不受行结构影响
    if G > 0:
        global_pos = np.fromiter(map(int, tokens[cursor:cursor + G]),
                                 dtype=np.int64, count=G)
        cursor += G
    else:
        global_pos = np.empty(0, dtype=np.int64)

    # Q、K、V 一次性读成一条扁平数组再 reshape：比逐行切片快得多，
    # 也避免 Python 层出现 3L 次 int() 转换。
    total = L * D
    flat = np.fromiter(map(int, tokens[cursor:cursor + 3 * total]),
                       dtype=np.int64, count=3 * total)
    Q = flat[:total].reshape(L, D)
    K = flat[total:2 * total].reshape(L, D)
    V = flat[2 * total:].reshape(L, D)
    # Q、K、V 是同一块扁平数组的三个视图，reshape/切片都不复制数据
    del tokens          # token 列表不再需要，提前释放

    # 是否全局位置摊成一个布尔掩码，逐行判断就是一次数组取值
    is_global = np.zeros(L, dtype=bool)
    is_global[global_pos] = True

    all_pos = np.arange(L)
    out = np.zeros((L, D), dtype=np.int64)   # 没有正分候选的行保持全 0

    def weighted_row(cand, scores):
        """丢掉非正分数，取分数最大的前 T 个，用它们加权求和得到一行 O[i]。"""
        keep = scores > 0                    # 0 也要丢掉，条件是严格大于 0
        if not keep.any():
            return None
        cand = cand[keep]
        scores = scores[keep]
        if scores.size > T:
            # lexsort 以最后一个键为主键：先按 -score（即分数降序），
            # 同分再按下标升序，与纯 Python 版的排序键 (-score, j) 完全等价。
            order = np.lexsort((cand, -scores))[:T]
            cand = cand[order]
            scores = scores[order]
        # 一次矩阵向量乘完成 O[i] = Σ score(i,j) * V[j]
        return scores @ V[cand]

    for i in range(L):
        if is_global[i]:
            # 全局位置的感受野是整条序列
            cand = all_pos
        else:
            lo = max(0, i - W)
            hi = min(L, i + W + 1)           # 左闭右开，右端记得 +1
            window = np.arange(lo, hi)
            # 只补窗口外的全局位置：窗口内的已经在上面算过一次，
            # 重复计入会把它的分数加两遍。
            extra = global_pos[(global_pos < lo) | (global_pos >= hi)]
            cand = np.concatenate((window, extra))

        row = weighted_row(cand, K[cand] @ Q[i])   # 候选行的点积一次算完
        if row is not None:
            out[i] = row

    sys.stdout.write("\n".join(" ".join(map(str, row)) for row in out.tolist()))


if __name__ == "__main__":
    solve()

---

示例推演

样例 1

$L=5$、$D=2$、$G=2$、$W=1$、$T=2$，全局位置 $P=\{0,3\}$。

三个矩阵各行的值：

下标	$Q$	$K$	$V$
0	[1, 0]	[2, 0]	[1, 1]
1	[1, 0]	[1, 0]	[2, 2]
2	[2, 0]	[3, 0]	[3, 3]
3	[0, 1]	[0, 2]	[4, 4]
4	[1, 1]	[-1, 1]	[5, 5]

注意 $K$ 的最后一行（下标 $4$）是 $[-1,1]$，它与 $Q[1]=[1,0]$ 的点积是 $-1$，会被丢掉。

逐行计算（每一行的 $C_i$ 按定义取窗口 $\cup$ 全局位置）：

$i$	是否全局	$C_i$	各候选分数	正分候选	取前 $T=2$	$O[i]$
0	是	$\{0,1,2,3,4\}$	$2,1,3,0,-1$	$j=0,1,2$	$2(3),\ 0(2)$	$[11,11]$
1	否	$\{0,1,2\}\cup\{0,3\}$	$2,1,3,0$	$j=0,1,2$	$2(3),\ 0(2)$	$[11,11]$
2	否	$\{1,2,3\}\cup\{0,3\}$	$4,2,6,0$	$j=0,1,2$	$2(6),\ 0(4)$	$[22,22]$
3	是	$\{0,1,2,3,4\}$	$0,0,0,2,1$	$j=3,4$	$3(2),\ 4(1)$	$[13,13]$
4	否	$\{3,4\}\cup\{0,3\}$	$2,2,0$	$j=0,3$	$0(2),\ 3(2)$	$[10,10]$

下面三行对应题目说明里的三次计算，可以对着表格核对：

$i=1$：窗口是 $\{0,1,2\}$，再并上全局位置 $0$ 和 $3$，得到 $C_1=\{0,1,2,3\}$。
分数分别是 $2,1,3,0$，丢掉 $0$ 之后有三个正分，按 $T=2$ 取 $j=2$ 与 $j=0$：

$$O[1]=3\times[3,3]+2\times[1,1]=[9,9]+[2,2]=[11,11].$$

$i=3$：$i$ 是全局位置，$C_3$ 是整条序列。分数只有 $j=3$ 的 $2$ 和 $j=4$ 的 $1$ 为正：

$$O[3]=2\times[4,4]+1\times[5,5]=[8,8]+[5,5]=[13,13].$$

$i=4$：窗口 $\{3,4\}$ 并上全局位置 $\{0,3\}$，去掉重复的 $3$ 得到 $C_4=\{0,3,4\}$。
分数为 $2,2,0$，$0$ 被丢掉；两个 $2$ 并列，先取下标小的 $0$，再取 $3$：

$$O[4]=2\times[1,1]+2\times[4,4]=[2,2]+[8,8]=[10,10].$$

另外两行也顺手验证一下：

• $i=0$ 是全局位置，$C_0$ 是全部位置，分数 $2,1,3,0,-1$，取 $j=2$（$3$）与 $j=0$（$2$），
  $O[0]=3\times[3,3]+2\times[1,1]=[11,11]$。
• $i=2$ 是普通位置，$C_2=\{0,1,2,3\}$，分数 $4,2,6,0$，取 $j=2$（$6$）与 $j=0$（$4$），
  $O[2]=6\times[3,3]+4\times[1,1]=[22,22]$。

五个输出行合起来正是样例输出：

11 11
11 11
22 22
13 13
10 10

样例 2

$L=3$、$D=1$、$G=0$、$W=0$、$T=1$，矩阵都是单列：
$Q=[1,2,-1]$、$K=[-2,3,1]$、$V=[10,20,30]$。

$G=0$ 表示没有全局位置，输入里也就没有那一行。$W=0$ 表示每个位置的窗口只有它自己。

$i$	$C_i$	分数	是否保留	$O[i]$
0	$\{0\}$	$1\times(-2)=-2$	否（非正）	$0$
1	$\{1\}$	$2\times 3=6$	是	$6\times20=120$
2	$\{2\}$	$(-1)\times1=-1$	否（非正）	$0$

输出 0 / 120 / 0，与样例一致。

---

复杂度分析

记每一行的候选个数为 $|C_i|\le G+2W+1$。

解法	时间复杂度	空间复杂度	说明
完整注意力矩阵	$O(L^2D+L^2\log L)$	$O(L^2)$	每行与全部 $L$ 个 key 打分，还要存 $L\times L$ 分数
按感受野模拟（排序取前 $T$）	$O\!\left(L\cdot\lvert C_i\rvert\cdot D+L\cdot\lvert C_i\rvert\log\lvert C_i\rvert\right)$	$O(LD)$	本题采用
按感受野模拟（只维护前 $T$ 大）	$O\!\left(L\cdot\lvert C_i\rvert\cdot D\right)$	$O(LD)$	用大小为 $T$ 的堆可去掉排序项

本题 $|C_i|\le 32+257=289$，$D\le64$，$L\le10^4$，
点积总数约 $2.9\times10^6$；排序每行最多 $289$ 个元素，和点积相比占比不大。

---

易错点总结

1. 把全局位置的候选也当成窗口

$i\in P$ 时 $C_i$ 是整条序列，不是窗口。全局位置存在的意义就是汇总远距离信息，
一旦改成窗口，它和普通位置就没有区别，窗口外的正分候选会被整个漏掉，
输出向量 $O[i]$ 随之偏小。

2. 忘记把全局位置并进普通位置的候选

只有窗口是不够的：$C_i=P\cup\{j:|i-j|\le W\}$。样例 1 的 $i=4$ 就是靠
全局位置 $0$ 拿到分数 $2$ 的。

3. 候选集合重复计数

全局位置可能正好落在窗口里。如果两条路径都把它加进去，它的分数会被加两次，
$O[i]$ 直接翻倍。代码里用 j < lo or j >= hi 过滤掉窗口内的全局位置。

4. 窗口边界

$i-W$ 可能小于 $0$，$i+W$ 可能大于 $L-1$。由于代码用左闭右开区间，
右边界要写 min(L, i + W + 1)，写成 min(L, i + W) 会少算最右边一个位置。

5. 分数等于 $0$ 也要丢掉

条件是 score > 0，不是 score >= 0。样例 1 里出现了好几个 $0$ 分，
保留它们会改变「取前 $T$ 个」的结果。

6. 并列时的顺序

分数相同时下标小的优先。代码用 (-score, j) 作为排序键，
升序排完就是「分数降序、下标升序」。如果只按分数排序，
下标顺序取决于排序算法是否稳定，结果不可靠。

7. 候选不足 $T$ 个时全部保留

题面说「若正分数候选个数不超过 $T$，则全部保留」，其实就是取前
$\min(T,\text{候选数})$ 个，代码里用 if len(candidates) > T 做截断即可，
不需要补齐到 $T$ 个。

8. 空集输出整行 $0$

一个正分候选都没有时，$O[i]$ 的每一维都是 $0$。注意是整行 $D$ 个 $0$，
不是输出一个 $0$；$D>1$ 时漏写会直接导致格式错误。

9. 这里没有 softmax，也不除以 $\sqrt D$

与标准的缩放点积注意力不同，本题的分数就是原始点积，权重就是分数本身，
既不归一化也不缩放。看到「注意力」就条件反射地加 softmax 会得到完全不同的结果。

10. $G=0$ 时输入少一行

全局位置那一行只在 $G>0$ 时存在。按行号硬编码解析会错位，
用 token 流顺序消费则不受影响。

---

扩展思考

• 全局位置的意义：让少数 token 拥有全局视野，是 Longformer、BigBird 等稀疏注意力的
  标准做法。它把 $O(L^2)$ 的注意力矩阵变成「带状窗口 + 若干整行/整列」，
  在保留远距离通路的同时把复杂度降到近似线性。
• 和纯滑动窗口的区别：P5337 的滑动窗口注意力只有一个窗口，位置之间要经过多层堆叠
  才能交换远距离信息；本题额外给了全局位置这条「捷径」，还额外做了正分过滤和 top-$T$ 截断。
• top-$T$ 的选法：本题 $T\le16$、$|C_i|\le289$，直接排序完全够用。
  如果 $T$ 和候选数都很大，可以改成维护一个大小为 $T$ 的小顶堆，
  把每行的 $|C_i|\log|C_i|$ 降到 $|C_i|\log T$。
• 候选集合更大时：若 $W$ 接近 $L$，窗口覆盖整条序列，本题就退化成完整注意力；
  此时再把点积循环换成向量化库（如 NumPy 的批量矩阵乘法）才有意义。
• 点积的常数：本题的瓶颈全在点积上。同样是 $O(L\cdot|C_i|\cdot D)$ 的写法，
  用 math.sumprod 之类的 C 实现能比纯 Python 双重循环快一个数量级，
  这也是代码里用 dot 封装一层的原因之一。

---

相关题目

• 滑动窗口注意力机制
• 模域双向循环注意力网络前向计算
• 滑动窗口