P5312. 模域双向循环注意力网络前向计算

场次：2026年8月26日 · 来源：CodeFun2000 P5312

题目描述

需要在模 $MOD$ 的整数域中，完成一个由多层双向循环神经网络和多头自注意力组成的网络的一次前向计算。

输入序列共有 $N$ 个位置，位置 $t$ 的原始输入向量记为 $x_t$，其维度为 $F$。网络先经过 $S$ 层双向 RNN 编码，再进行包含 $A$ 个注意力头的多头自注意力计算，之后通过残差连接，最后使用线性输出层得到 $C$ 维向量 $y_t$。

所有向量均视为行向量，矩阵乘法均采用“行向量乘矩阵”的方式。所有输入参数均为满足 $0\le x<MOD$ 的非负整数。每次完成向量加法、点积或矩阵乘法后，都需要将结果对 $MOD$ 取模。

1. 多层双向 RNN 编码

对于双向 RNN 的第 $l$ 层，位置 $t$ 的输入记为 $a_t^{(l)}$。第一层直接使用原始输入，即 $a_t^{(1)}=x_t$。

每一层的输入维度记作 $d_l$，其中 $d_l=\begin{cases}F,&l=1\\2R,&l>1\end{cases}$。

这里，$R$ 是单个方向的隐藏状态维度。因此，一层双向 RNN 在每个位置都会分别产生一个 $R$ 维正向状态和一个 $R$ 维反向状态。

第 $l$ 层正向计算所使用的参数为 $U_f^{(l)}$、$V_f^{(l)}$、$c_f^{(l)}$，其规模依次为 $d_l\times R$、$R\times R$、$1\times R$。该方向的给定初始状态为 $p_0^{(l)}$，规模为 $1\times R$。

正向状态从位置 $1$ 开始，按照 $1,2,\dots,N$ 的顺序计算。对于每个 $t$，有 $p_t^{(l)}=\left(a_t^{(l)}U_f^{(l)}+p_{t-1}^{(l)}V_f^{(l)}+c_f^{(l)}\right)\bmod MOD$。

同一层的反向计算使用 $U_b^{(l)}$、$V_b^{(l)}$、$c_b^{(l)}$，规模分别为 $d_l\times R$、$R\times R$、$1\times R$。反向方向的给定初始状态为 $r_{N+1}^{(l)}$，其规模为 $1\times R$。

反向状态的计算次序与正向相反，从位置 $N$ 开始依次处理到位置 $1$。对于每个 $t$，有 $r_t^{(l)}=\left(a_t^{(l)}U_b^{(l)}+r_{t+1}^{(l)}V_b^{(l)}+c_b^{(l)}\right)\bmod MOD$。

当第 $l$ 层的两个方向均计算完成后，把同一位置上的正向状态放在前、反向状态放在后进行拼接。若 $l<S$，该结果作为下一层对应位置的输入，即 $a_t^{(l+1)}=\left[p_t^{(l)}\parallel r_t^{(l)}\right]$。

因此，从第 $2$ 层开始，每个 $a_t^{(l)}$ 的维度均为 $2R$。

经过全部 $S$ 层双向 RNN 后，取最后一层在位置 $t$ 的两个方向状态，并定义 $e_t=\left[p_t^{(S)}\parallel r_t^{(S)}\right]$。

此时 $e_t\in\mathbb{Z}_{MOD}^{2R}$，它将作为后续多头自注意力部分的输入。

2. 多头自注意力

多头自注意力共有 $A$ 个注意力头，注意力头编号为 $h=0,1,\dots,A-1$，每个头的维度为 $Q$。

本题使用未归一化的点积注意力，不进行缩放、不使用 softmax，也不存在 mask。

对于任意注意力头 $h$，均给定三个规模为 $2R\times Q$ 的矩阵 $Q_h$、$K_h$、$V_h$。在位置 $t$，分别计算 $q_t^{(h)}=e_tQ_h\bmod MOD$、$k_t^{(h)}=e_tK_h\bmod MOD$、$v_t^{(h)}=e_tV_h\bmod MOD$。

上述三个向量均为 $Q$ 维，即 $q_t^{(h)},k_t^{(h)},v_t^{(h)}\in\mathbb{Z}_{MOD}^{Q}$。

接下来，对于位置 $t$ 和位置 $s$，利用 $q_t^{(h)}$ 与 $k_s^{(h)}$ 的点积计算第 $h$ 个注意力头的注意力分数：

$\beta_{t,s}^{(h)}=\left(\sum_{j=0}^{Q-1}q_t^{(h)}[j]\cdot k_s^{(h)}[j]\right)\bmod MOD$。

得到位置 $t$ 对所有位置的注意力分数后，对每个分量 $j=0,1,\dots,Q-1$ 计算：

$c_t^{(h)}[j]=\left(\sum_{s=1}^{N}\beta_{t,s}^{(h)}\cdot v_s^{(h)}[j]\right)\bmod MOD$。

由此得到第 $h$ 个注意力头在位置 $t$ 的 $Q$ 维输出 $c_t^{(h)}$。

对于同一个位置 $t$，将各注意力头的结果按照编号 $0,1,\dots,A-1$ 的顺序拼接，得到 $c_t=\left[c_t^{(0)}\parallel c_t^{(1)}\parallel\dots\parallel c_t^{(A-1)}\right]$。

拼接后的向量满足 $c_t\in\mathbb{Z}_{MOD}^{AQ}$。

3. 残差连接

多头自注意力产生的 $AQ$ 维向量 $c_t$ 需要通过投影矩阵 $W_p$ 映射到 $2R$ 维。

其中，$W_p$ 的规模为 $AQ\times 2R$，偏置 $c_p$ 的规模为 $1\times 2R$。

投影结果与同一位置的编码结果 $e_t$ 进行残差相加，得到 $u_t=\left(c_tW_p+c_p+e_t\right)\bmod MOD$。

4. 输出层

残差连接后的 $u_t$ 通过最后的线性层得到位置 $t$ 的输出。

输出层包含规模为 $2R\times C$ 的矩阵 $W_y$ 和规模为 $1\times C$ 的偏置 $c_y$，计算方式为 $y_t=\left(u_tW_y+c_y\right)\bmod MOD$。

最终得到 $y_t\in\mathbb{Z}_{MOD}^{C}$。需要计算输入序列全部 $N$ 个位置对应的最终输出向量。

输入格式

第一行给出 $8$ 个整数 $N$、$F$、$R$、$C$、$S$、$A$、$Q$、$MOD$，分别表示序列长度、输入维度、单向隐藏状态维度、输出维度、双向 RNN 层数、注意力头数量、每个注意力头的维度以及模数。

数据范围为：

$1\le N\le 10,\ 1\le F\le 5,\ 1\le R\le 5$。

$1\le C\le 5,\ 1\le S\le 3,\ 1\le A\le 4,\ 1\le Q\le 4$。

$2\le MOD\le 97$。

随后给出原始输入序列，共 $N$ 行。第 $t$ 行包含 $F$ 个整数，表示位置 $t$ 的向量 $x_t$。

原始输入之后是 $S$ 层双向 RNN 的参数，各层按照 $l=1,2,\dots,S$ 的顺序给出。对于当前层 $l$，输入维度为 $d_l=\begin{cases}F,&l=1\\2R,&l>1\end{cases}$。

每一层内部的参数输入顺序固定如下：

1. $U_f^{(l)}$：输入 $d_l$ 行，每行包含 $R$ 个整数；
2. $V_f^{(l)}$：输入 $R$ 行，每行包含 $R$ 个整数；
3. $c_f^{(l)}$：输入 $1$ 行，包含 $R$ 个整数；
4. $U_b^{(l)}$：输入 $d_l$ 行，每行包含 $R$ 个整数；
5. $V_b^{(l)}$：输入 $R$ 行，每行包含 $R$ 个整数；
6. $c_b^{(l)}$：输入 $1$ 行，包含 $R$ 个整数；
7. $p_0^{(l)}$：输入 $1$ 行，包含 $R$ 个整数；
8. $r_{N+1}^{(l)}$：输入 $1$ 行，包含 $R$ 个整数。

全部 $S$ 层双向 RNN 参数之后，依次给出 $A$ 个注意力头的参数。注意力头按照 $h=0,1,\dots,A-1$ 的编号顺序输入，每个头依次包含：

1. $Q_h$：输入 $2R$ 行，每行包含 $Q$ 个整数；
2. $K_h$：输入 $2R$ 行，每行包含 $Q$ 个整数；
3. $V_h$：输入 $2R$ 行，每行包含 $Q$ 个整数。

所有注意力头的参数结束后，剩余参数按照 $W_p$、$c_p$、$W_y$、$c_y$ 的顺序给出：

1. $W_p$：输入 $AQ$ 行，每行包含 $2R$ 个整数；
2. $c_p$：输入 $1$ 行，包含 $2R$ 个整数；
3. $W_y$：输入 $2R$ 行，每行包含 $C$ 个整数；
4. $c_y$：输入 $1$ 行，包含 $C$ 个整数。

所有输入矩阵均按照从上到下、逐行的顺序给出。

输出格式

按照位置 $1,2,\dots,N$ 的顺序输出计算结果，共输出 $N$ 行。

第 $t$ 行包含 $C$ 个非负整数，依次表示最终输出向量 $y_t$ 的各个分量，相邻整数之间使用一个空格分隔。

示例

示例 1

输入：

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

输出：

1

说明：

由于 $N=1$，只需计算位置 1。第 1 层正向隐藏状态为 $p_1^{(1)} = (2 \cdot 3 + 4 \cdot 4 + 1) \bmod 5 = 23 \bmod 5 = 3$。 反向隐藏状态为 $r_1^{(1)} = (2 \cdot 2 + 0 \cdot 1 + 3) \bmod 5 = 7 \bmod 5 = 2$。 因此编码器输出 $e_1=[3,2]$。

注意力头中： $q_1 = e_1 Q_0 = 3 \cdot 1 + 2 \cdot 0 = 3$， $k_1 = e_1 K_0 = 3 \cdot 0 + 2 \cdot 1 = 2$， $v_1 = e_1 V_0 = (3 \cdot 4 + 2 \cdot 3) \bmod 5 = 18 \bmod 5 = 3$。 注意力分数 $\beta_{1,1} = q_1 k_1 \bmod 5 = 6 \bmod 5 = 1$，因此注意力输出 $c_1 = 1 \cdot 3 \bmod 5 = 3$。

投影层中 $c_1 W_p = [3 \cdot 1, 3 \cdot 4] \bmod 5 = [3,2]$。 残差连接后 $u_1 = ([3,2]+[2,1]+[3,2]) \bmod 5 = [3,0]$。 输出层计算 $y_1 = (3 \cdot 2 + 0 \cdot 1 + 0) \bmod 5 = 6 \bmod 5 = 1$。

最终输出为 1。

示例 2

输入：

2 1 1 2 1 1 1 5
1
2
1
0
0
1
0
0
0
0
1
0
0
1
1
0
1 1
0 0
1 0
0 1
0 0

输出：

1 1
2 2

说明：

第 1 层中正向和反向都只使用当前输入，因此 $p_t^{(1)} = x_t$，$r_t^{(1)} = x_t$。 于是位置 1 的编码器输出为 $e_1=[1,1]$，位置 2 的编码器输出为 $e_2=[2,2]$。

注意力头参数将 $q_t$ 取为 $e_t$ 的第 1 个分量，$k_t$ 取为 $e_t$ 的第 2 个分量，$v_t$ 取为 $e_t$ 的第 1 个分量。 所以 $q_1=1$，$q_2=2$，$k_1=1$，$k_2=2$，$v_1=1$，$v_2=2$。

对于位置 1，注意力分数为 $\beta_{1,1}=1 \cdot 1 \bmod 5=1$， $\beta_{1,2}=1 \cdot 2 \bmod 5=2$。 注意力输出为 $c_1 = (1 \cdot 1 + 2 \cdot 2) \bmod 5 = 5 \bmod 5 = 0$。

对于位置 2，注意力分数为 $\beta_{2,1}=2 \cdot 1 \bmod 5=2$， $\beta_{2,2}=2 \cdot 2 \bmod 5=4$。 注意力输出为 $c_2 = (2 \cdot 1 + 4 \cdot 2) \bmod 5 = 10 \bmod 5 = 0$。

所以两个位置的注意力输出均为 0。残差后 $u_t = e_t$，输出层使用单位矩阵，因此最终输出为位置 1 的 $[1,1]$ 和位置 2 的 $[2,2]$。

示例 3

输入：

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

输出：

2

说明：

由于 $N=1$，仍只需计算位置 1。第 1 层正向隐藏状态为 $p_1^{(1)} = (2 \cdot 3 + 4 \cdot 4 + 1) \bmod 5 = 3$。 反向隐藏状态为 $r_1^{(1)} = (2 \cdot 2 + 0 \cdot 1 + 3) \bmod 5 = 2$。 因此编码器输出 $e_1=[3,2]$。

第 0 个头中： $q_1^{(0)} = 3 \cdot 1 + 2 \cdot 0 = 3$， $k_1^{(0)} = 3 \cdot 0 + 2 \cdot 1 = 2$， $v_1^{(0)} = (3 \cdot 4 + 2 \cdot 3) \bmod 5 = 3$。 注意力分数 $\beta_{1,1}^{(0)} = 3 \cdot 2 \bmod 5 = 6 \bmod 5 = 1$， 因此该头输出 $c_1^{(0)} = 1 \cdot 3 = 3$。

第 1 个头中： $q_1^{(1)} = 3 \cdot 0 + 2 \cdot 1 = 2$， $k_1^{(1)} = (3 \cdot 1 + 2 \cdot 2) \bmod 5 = 7 \bmod 5 = 2$， $v_1^{(1)} = (3 \cdot 1 + 2 \cdot 4) \bmod 5 = 11 \bmod 5 = 1$。 注意力分数 $\beta_{1,1}^{(1)} = 2 \cdot 2 \bmod 5 = 4$， 因此该头输出 $c_1^{(1)} = 4 \cdot 1 = 4$。

拼接后注意力输出为 $c_1=[3,4]$。投影矩阵是单位矩阵，所以残差后 $u_1 = ([3,4]+[0,0]+[3,2]) \bmod 5 = [1,1]$。 输出层计算 $y_1 = (1 \cdot 1 + 1 \cdot 1 + 0) \bmod 5 = 2$。

最终输出为 2。

---

解题思路

第一步：理解问题本质

题目要求按照给定公式，完成一个四段式网络的前向计算：

1. 多层双向 RNN：每一层分别从左到右、从右到左计算隐藏状态，再把两个方向拼接；
2. 多头自注意力：每个头独立计算 $q,k,v$，用点积得到注意力分数，再对 $v$ 加权求和；
3. 残差连接：把注意力输出投影回 $2R$ 维，并与编码结果 $e_t$ 相加；
4. 输出层：经过线性变换得到 $C$ 维输出。

所有运算都在模 $MOD$ 的整数域中进行。题目数据范围很小（$N\le10$、$R\le5$、$Q\le4$ 等），因此不需要复杂的优化，关键是严格按照题目给定的顺序和维度实现。

第二步：朴素模拟

最直接的做法是把每个模块拆成独立的矩阵运算：

• RNN 的每一步都写三重循环；
• 每个注意力头都重新计算 $q_t,k_s,v_s$；
• 每一步运算后都取模。

这种做法是正确的，但代码容易写乱，尤其是参数输入顺序和多层 RNN 的维度变化。题目本身没有“更快的渐进算法”，因为所有矩阵都很小；真正的难点是不能把公式抄错。

第三步：工程化实现

把重复的矩阵运算抽成两个辅助函数：

# 行向量乘矩阵：vec 的长度必须等于 mat 的行数，结果的第 j 个分量是 vec 与 mat 第 j 列的点积
def vec_mat(vec, mat):
    """行向量乘矩阵"""
    # 输出向量的维度等于矩阵的列数，取第一行的长度即可（输入保证每行列数相同）
    cols = len(mat[0])
    return [
        # 每个分量独立求和后立刻取模：题目要求每次点积都取模，而对每个分量分别取模结果一致，也最不容易漏
        sum(vec[i] * mat[i][j] for i in range(len(vec))) % MOD
        for j in range(cols)
    ]

这样 RNN、注意力、投影和输出层都可以复用同一个函数，减少抄错维度的机会。

RNN 部分按层处理，只保留当前层的输入 a：

• 第一层的输入维度是 $F$；
• 后续层的输入是上一层正向、反向状态的拼接，维度是 $2R$；
• 正向递推使用 $p_0$，反向递推使用 $r_{N+1}$；
• 每层计算完成后，用 p[t] + r[t] 作为下一层的输入。

注意力部分按头处理：

• 先对每个位置计算 $q_t=e_tQ_h$、$k_t=e_tK_h$、$v_t=e_tV_h$；
• 再计算 $\beta_{t,s}=q_t\cdot k_s$；
• 最后计算 $c_t[j]=\sum_s\beta_{t,s}v_s[j]$；
• 所有头算完后按编号顺序拼接。

第四步：完整流程与下标约定

题目公式中的位置是 $1$ 起始的，代码中统一使用 $0$ 起始的下标：

• 正向递推：p[0] 从 p_0 开始，p[t] 由 p[t-1] 推出；
• 反向递推：r[N-1] 从 r_{N+1} 开始，r[t] 由 r[t+1] 推出。

拼接顺序必须严格保持：

$$e_t=\left[p_t^{(S)}\parallel r_t^{(S)}\right]$$

即正向状态在前，反向状态在后。注意力头的拼接顺序同理，必须是 $h=0,1,\dots,A-1$。

残差连接使用的是编码结果 $e_t$，不是最后一层的输入 $a_t$ 的某个中间状态：

$$u_t=\left(c_tW_p+c_p+e_t\right)\bmod MOD$$

最后输出：

$$y_t=\left(u_tW_y+c_y\right)\bmod MOD$$

---

完整代码实现

算法可视化 · 逐帧回放这道题的代码133 帧 · 录制输入

import sys


def solve():
    # 本题所有输入都是整数，一次性读入后按顺序取用即可。
    data = list(map(int, sys.stdin.buffer.read().split()))
    if not data:
        return

    # 头部 8 个整数：序列长度、输入维度、隐藏维度、输出维度、
    # RNN 层数、注意力头数、每个头的维度、模数。
    pos = 0
    N, F, R, C, S, A, Q, MOD = data[pos:pos + 8]
    pos += 8

    def read_matrix(rows, cols):
        """从当前位置读取 rows x cols 的整数矩阵。"""
        nonlocal pos
        mat = []
        for _ in range(rows):
            # 输入保证 0 <= x < MOD，这里仍然统一取模，避免手误。
            row = data[pos:pos + cols]
            pos += cols
            mat.append([x % MOD for x in row])
        return mat

    def read_vector(size):
        """从当前位置读取一个长度为 size 的整数向量。"""
        return read_matrix(1, size)[0]

    def vec_mat(vec, mat):
        """行向量乘矩阵：result[j] = sum_i vec[i] * mat[i][j]。"""
        cols = len(mat[0])
        return [
            sum(vec[i] * mat[i][j] for i in range(len(vec))) % MOD
            for j in range(cols)
        ]

    # a 表示“当前 RNN 层的输入序列”：
    # - 第 1 层：a 是原始输入，形状 N x F；
    # - 第 2 层开始：a 是上一层正向/反向状态的拼接，形状 N x 2R。
    a = read_matrix(N, F)

    # 逐层计算双向 RNN
    for layer in range(S):
        # 第 1 层的输入维度是 F，之后每层都是 2R。
        d = F if layer == 0 else 2 * R

        # 正向参数：U_f (d x R)、V_f (R x R)、c_f (1 x R)
        U_f = read_matrix(d, R)
        V_f = read_matrix(R, R)
        c_f = read_vector(R)

        # 反向参数：U_b (d x R)、V_b (R x R)、c_b (1 x R)
        U_b = read_matrix(d, R)
        V_b = read_matrix(R, R)
        c_b = read_vector(R)

        # 正向初始状态 p_0 和反向初始状态 r_{N+1}
        p_prev = read_vector(R)
        r_next = read_vector(R)

        # 正向递推：p_t = a_t U_f + p_{t-1} V_f + c_f
        p = [None] * N
        prev = p_prev
        for t in range(N):
            cur = [
                (sum(a[t][i] * U_f[i][j] for i in range(d))
                 + sum(prev[i] * V_f[i][j] for i in range(R))
                 + c_f[j]) % MOD
                for j in range(R)
            ]
            p[t] = cur
            prev = cur

        # 反向递推：r_t = a_t U_b + r_{t+1} V_b + c_b
        # 从最后一个位置 N-1 开始，向前推到位置 0。
        r = [None] * N
        nxt = r_next
        for t in range(N - 1, -1, -1):
            cur = [
                (sum(a[t][i] * U_b[i][j] for i in range(d))
                 + sum(nxt[i] * V_b[i][j] for i in range(R))
                 + c_b[j]) % MOD
                for j in range(R)
            ]
            r[t] = cur
            nxt = cur

        # 正向状态在前、反向状态在后，拼接成下一层的输入。
        # 拼接后每个位置的维度是 R + R = 2R。
        a = [p[t] + r[t] for t in range(N)]

    # 全部 RNN 层结束后，a 就是编码结果 e_t。
    e = a

    # 多头自注意力：每个头都有自己的一套 Q_h、K_h、V_h。
    head_outputs = []
    for _ in range(A):
        # 每个矩阵的形状都是 2R x Q，因为 e_t 是 2R 维。
        Q_h = read_matrix(2 * R, Q)
        K_h = read_matrix(2 * R, Q)
        V_h = read_matrix(2 * R, Q)

        # 对每个位置 t，分别计算 Q、K、V 投影。
        q = [vec_mat(e[t], Q_h) for t in range(N)]
        k = [vec_mat(e[t], K_h) for t in range(N)]
        v = [vec_mat(e[t], V_h) for t in range(N)]

        # beta[t][s] = q_t · k_s，表示位置 t 对位置 s 的注意力分数。
        beta = [
            [
                sum(q[t][j] * k[s][j] for j in range(Q)) % MOD
                for s in range(N)
            ]
            for t in range(N)
        ]

        # c_t[j] = sum_s beta[t][s] * v_s[j]，
        # 即用注意力分数对所有位置的 V 做加权求和。
        c = [
            [
                sum(beta[t][s] * v[s][j] for s in range(N)) % MOD
                for j in range(Q)
            ]
            for t in range(N)
        ]
        head_outputs.append(c)

    # 按头编号拼接：c_t = [c_t^(0) || c_t^(1) || ... || c_t^(A-1)]。
    # 拼接后每个位置是 AQ 维。
    c = [
        [head_outputs[h][t][j] for h in range(A) for j in range(Q)]
        for t in range(N)
    ]

    # 投影 + 残差：
    # W_p 是 AQ x 2R，c_p 是 2R 维，e_t 也是 2R 维。
    # u_t = c_t W_p + c_p + e_t。
    W_p = read_matrix(A * Q, 2 * R)
    c_p = read_vector(2 * R)
    u = [
        [
            (sum(c[t][i] * W_p[i][j] for i in range(A * Q))
             + c_p[j] + e[t][j]) % MOD
            for j in range(2 * R)
        ]
        for t in range(N)
    ]

    # 输出层：W_y 是 2R x C，c_y 是 C 维。
    # y_t = u_t W_y + c_y。
    W_y = read_matrix(2 * R, C)
    c_y = read_vector(C)
    y = [
        [
            (sum(u[t][i] * W_y[i][j] for i in range(2 * R)) + c_y[j]) % MOD
            for j in range(C)
        ]
        for t in range(N)
    ]

    for row in y:
        print(*row)


if __name__ == "__main__":
    solve()

NumPy 版本

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

思路与上面的纯 Python 版完全一致，区别在于把「逐分量」的循环换成矩阵乘法：RNN 的输入投影 $a_tU$、注意力的 $q/k/v$、注意力分数与加权求和、投影层与输出层，全部一次算完，并且每次乘法后立刻 % MOD。元素都小于 $MOD\le 97$，累加上限是 $10\times97\times97<10^5$，int64 有充足余量，不存在溢出。

import sys

import numpy as np


def solve():
    """按题面公式做一次前向计算；所有矩阵乘法用 int64 并在每次乘完后取模。"""
    tokens = sys.stdin.buffer.read().split()
    if not tokens:
        return

    # 头部 8 个整数：序列长度、输入维度、隐藏维度、输出维度、
    # RNN 层数、注意力头数、每个头的维度、模数。
    N, F, R, C, S, A, Q, MOD = (int(x) for x in tokens[:8])
    pos = 8

    def read_matrix(rows, cols):
        """从当前位置读取 rows x cols 的整数矩阵，形状为 (rows, cols)。"""
        nonlocal pos
        mat = np.array(tokens[pos:pos + rows * cols], dtype=np.int64).reshape(rows, cols)
        pos += rows * cols
        # 输入保证 0 <= x < MOD，这里统一取模，避免手误。
        return mat % MOD

    def read_vector(size):
        """从当前位置读取一个长度为 size 的整数向量。"""
        return read_matrix(1, size)[0]

    # a 表示“当前 RNN 层的输入序列”，每一行是一个位置：
    # - 第 1 层：a 是原始输入，形状 N x F；
    # - 第 2 层开始：a 是上一层正向/反向状态的拼接，形状 N x 2R。
    a = read_matrix(N, F)

    # 逐层计算双向 RNN。这里 int64 足够宽：元素都小于 MOD <= 97，
    # 一次矩阵乘法的累加上限是 10 * 97 * 97 < 1e5，远达不到 2^63；
    # 每次乘法后立刻取模，量级就始终被压在 MOD 以内。
    for layer in range(S):
        # 第 1 层的输入维度是 F，之后每层都是 2R。
        d = F if layer == 0 else 2 * R

        # 正向参数：U_f (d x R)、V_f (R x R)、c_f (1 x R)
        U_f = read_matrix(d, R)
        V_f = read_matrix(R, R)
        c_f = read_vector(R)

        # 反向参数：U_b (d x R)、V_b (R x R)、c_b (1 x R)
        U_b = read_matrix(d, R)
        V_b = read_matrix(R, R)
        c_b = read_vector(R)

        # 正向初始状态 p_0 和反向初始状态 r_{N+1}
        p_prev = read_vector(R)
        r_next = read_vector(R)

        # 输入投影与递推无关，可以一次性算完所有位置：in_f[t] = a_t U_f。
        in_f = (a @ U_f) % MOD
        in_b = (a @ U_b) % MOD

        # 正向递推：p_t = a_t U_f + p_{t-1} V_f + c_f。
        # 这一步有前后依赖，只能按位置顺序推进，但每一步都是一次向量化乘法。
        p = np.empty((N, R), dtype=np.int64)
        prev = p_prev
        for t in range(N):
            cur = (in_f[t] + prev @ V_f + c_f) % MOD
            p[t] = cur
            prev = cur

        # 反向递推：r_t = a_t U_b + r_{t+1} V_b + c_b，从位置 N-1 向 0 推进。
        r = np.empty((N, R), dtype=np.int64)
        nxt = r_next
        for t in range(N - 1, -1, -1):
            cur = (in_b[t] + nxt @ V_b + c_b) % MOD
            r[t] = cur
            nxt = cur

        # 正向状态在前、反向状态在后，拼接成下一层的输入（每行 2R 维）。
        a = np.concatenate([p, r], axis=1)

    # 全部 RNN 层结束后，a 就是编码结果 e_t，形状 N x 2R。
    e = a

    # 多头自注意力：每个头一套 Q_h、K_h、V_h，都是 2R x Q。
    head_outputs = []
    for _ in range(A):
        Q_h = read_matrix(2 * R, Q)
        K_h = read_matrix(2 * R, Q)
        V_h = read_matrix(2 * R, Q)

        # 一次矩阵乘法算出所有位置的 q/k/v，形状都是 N x Q。
        q = (e @ Q_h) % MOD
        k = (e @ K_h) % MOD
        v = (e @ V_h) % MOD

        # beta[t][s] = q_t · k_s：q 乘 k 的转置，一次得到 N x N 的注意力分数。
        beta = (q @ k.T) % MOD

        # c_t = sum_s beta[t][s] * v_s：用注意力分数对所有位置的 V 加权求和。
        c = (beta @ v) % MOD
        head_outputs.append(c)

    # 按头编号拼接：c_t = [c_t^(0) || c_t^(1) || ... || c_t^(A-1)]，每行 AQ 维。
    c = np.concatenate(head_outputs, axis=1)

    # 投影 + 残差：u_t = (c_t W_p + c_p + e_t) mod MOD。
    W_p = read_matrix(A * Q, 2 * R)
    c_p = read_vector(2 * R)
    u = (c @ W_p + c_p + e) % MOD

    # 输出层：y_t = (u_t W_y + c_y) mod MOD，形状 N x C。
    W_y = read_matrix(2 * R, C)
    c_y = read_vector(C)
    y = (u @ W_y + c_y) % MOD

    for row in y.tolist():
        print(*row)


if __name__ == "__main__":
    solve()

---

示例推演

样例 1

输入规模为 $N=1,F=1,R=1,C=1,S=1,A=1,Q=1,MOD=5$。

双向 RNN：

$$p_1^{(1)}=(2\cdot3+4\cdot4+1)\bmod5=23\bmod5=3$$ $$r_1^{(1)}=(2\cdot2+0\cdot1+3)\bmod5=7\bmod5=2$$

所以：

$$e_1=[p_1\parallel r_1]=[3,2]$$

注意力：

$$q_1=e_1Q_0=3\cdot1+2\cdot0=3$$ $$k_1=e_1K_0=3\cdot0+2\cdot1=2$$ $$v_1=e_1V_0=3\cdot4+2\cdot3=18\bmod5=3$$

注意力分数：

$$\beta_{1,1}=q_1k_1\bmod5=6\bmod5=1$$

注意力输出：

$$c_1=\beta_{1,1}v_1=1\cdot3=3$$

残差与输出：

$$c_1W_p=[3,2],\qquad c_p=[2,1],\qquad e_1=[3,2]$$ $$u_1=([3,2]+[2,1]+[3,2])\bmod5=[3,0]$$ $$y_1=(3\cdot2+0\cdot1+0)\bmod5=1$$

最终输出 1。

样例 3：多头拼接

样例 3 在样例 1 的基础上把注意力头数量改成 $A=2$。

• 第 0 个头输出 $c_1^{(0)}=3$；
• 第 1 个头输出 $c_1^{(1)}=4$。

按头编号拼接：

$$c_1=[c_1^{(0)}\parallel c_1^{(1)}]=[3,4]$$

投影矩阵是单位矩阵，所以：

$$u_1=([3,4]+[0,0]+[3,2])\bmod5=[1,1]$$

输出层：

$$y_1=(1\cdot1+1\cdot1+0)\bmod5=2$$

最终输出 2。这个例子说明：头的拼接顺序和残差使用的是编码结果 $e_t$，两处都不能混用。

---

复杂度分析

模块	时间复杂度	空间复杂度	说明
双向 RNN	$O(S\cdot N\cdot(dR+R^2))$	$O(NR)$	每层正向、反向各扫一遍
多头注意力	$O(A\cdot N^2\cdot Q)$	$O(ANQ)$	主要开销在 $\beta$ 和加权求和
残差投影	$O(N\cdot AQ\cdot R)$	$O(NR)$	把 $AQ$ 维映射回 $2R$ 维
输出层	$O(N\cdot R\cdot C)$	$O(NC)$	线性变换
总计	$O(S\cdot N\cdot R(F+R)+A\cdot N^2Q+NAQR+NRC)$	$O(NR+ANQ)$	本题数据范围很小

由于 $N\le10$、$S\le3$、$A\le4$、$Q\le4$，即使直接按公式模拟，计算量也非常小。

---

易错点总结

1. RNN 参数顺序不能乱

每一层 RNN 的参数顺序固定为：

$$U_f,\ V_f,\ c_f,\ U_b,\ V_b,\ c_b,\ p_0,\ r_{N+1}$$

其中 $U_f$ 和 $U_b$ 的行数是当前层的输入维度 $d_l$，$V_f$ 和 $V_b$ 是 $R\times R$。读错一个矩阵，后面的结果会全部错位。

2. 第一层和后续层的输入维度不同

$$d_l=\begin{cases}F,&l=1\\2R,&l>1\end{cases}$$

第一层直接使用原始输入 $x_t$，从第二层开始使用上一层拼接后的 $2R$ 维状态。

3. 正向和反向的初始状态不同

• 正向从 $p_0$ 开始，按 $1,2,\dots,N$ 递推；
• 反向从 $r_{N+1}$ 开始，按 $N,N-1,\dots,1$ 递推。

反向递推时，当前位置使用的是后一个位置的状态 r[t+1]，不能写成 r[t-1]。

4. 拼接顺序固定

编码结果是：

$$e_t=[p_t\parallel r_t]$$

正向在前，反向在后。注意力头的结果同样按 $0,1,\dots,A-1$ 的顺序拼接。顺序反了，投影矩阵的列就对不上。

5. 残差加的是 $e_t$

残差连接是：

$$u_t=(c_tW_p+c_p+e_t)\bmod MOD$$

这里加的是编码器输出 $e_t$，不是最后一层的输入 $a_t$，也不是注意力输出 $c_t$。

6. 每一步都要取模

题目要求向量加法、点积和矩阵乘法后都对 $MOD$ 取模。Python 的 % 对正数结果非负，可以直接使用。所有输入本身也满足 $0\le x<MOD$，但读取时统一 % MOD 更稳妥。

7. 输出是整数，不是浮点数

本题所有运算都在模域中进行，输出的是 $0$ 到 $MOD-1$ 之间的整数，不需要保留小数，也不要用浮点运算代替整数取模。

---

扩展思考

• 模运算与矩阵乘法：由于加法和乘法对模运算封闭，可以先累加再取模，结果与每一步取模相同。本题为了代码清晰，在每个分量计算完成后统一取模。
• 矩阵乘法的顺序：题目明确所有向量都是行向量，计算顺序是“行向量乘矩阵”。如果改成列向量，所有矩阵都要转置，公式会完全不同。
• 复杂度瓶颈：本题的瓶颈是多头注意力的 $O(A N^2 Q)$。当 $N$ 很大时，可以用滑动窗口注意力把 $N^2$ 降到 $Nw$，这正是下一题 P5337 的内容。
• 工程实现：如果维度更大，可以把 vec_mat 改成先转置矩阵、再按行做点积，或者使用 NumPy 的批量矩阵乘法；本题数据范围小，不需要这些优化。

---

相关题目

• 滑动窗口注意力机制