这道题的正解是线段树维护可合并区间状态,但我们用朴素模拟合并也能在考试时拿到一定的分数。每次区间查询从 l 到 r 依次合并,单点更新直接改数组;B、Q 较小时完全够用,大数据会超时。
#code-switcher
import math
给定一个由 B 个序列块组成的有序序列。第 i 个块保存状态 (mxi,wdi,svi)。在接下来的 Q 次操作中,某个块的状态可能被重新赋值,也可能需要求出一段连续区间合并之后对应的注意力结果。
这些状态来自 Transformer 的 Attention 计算。FlashAttention 为降低长序列带来的内存开销,会将计算拆分为多个块并以流式方式维护局部信息。对于任意一个块,其状态由 (mx,wd,sv) 三部分构成,含义如下:
在该块涉及的所有 q⋅kj 中,最大值记录为 mx;
以 mx 为基准进行重缩放后,所有权重的总和记录为 wd,即
wd=∑exp(q⋅kj−mx);对 vj 使用相同的重缩放权重进行加权,所得总和记录为 sv,即
sv=∑exp(q⋅kj−mx)⋅vj.因此,一个状态所表示的注意力结果可以直接由
attn=wdsv得到。
当两个相邻块已经分别汇总为
A=(mxA,wdA,svA)和
B=(mxB,wdB,svB)时,需要先统一二者的缩放基准,再得到合并状态 C=(mxC,wdC,svC)。具体计算规则为
mxCwdCsvC=max(mxA,mxB),=wdA⋅emxA−mxC+wdB⋅emxB−mxC,=svA⋅emxA−mxC+svB⋅emxB−mxC.这一合并运算满足结合律。也就是说,在合并一段连续的块时,可以采用任意括号顺序进行计算,所得最终状态不会发生变化。
现在需要处理以下两种操作:
单点更新:给定位置 i 以及新的 (mx,wd,sv),用该三元组替换第 i 个块原有的状态;
区间查询:给定 [l,r],将第 l 个块到第 r 个块依次合并。若最终状态为 (mxC,wdC,svC),则需要输出
wdCsvC.第一行输入两个整数 B,Q,分别表示块的总数和操作的总数。
随后有 B 行初始状态信息。第 i 行给出三个实数 mxi wdi svi,表示第 i 个块的状态,其中 1≤i≤B。
再接下来输入 Q 行操作,每行采用以下两种格式之一:
1 i mx wd sv:执行单点更新,将第 i 个块修改为状态 (mx,wd,sv),本次操作无需输出;2 l r:执行区间查询,求第 l 个块至第 r 个块合并后的注意力结果。所有输入满足:
合并过程中需要进行指数计算,请使用双精度浮点(double / float),指数函数使用标准库 exp()。
每遇到一次区间查询操作,输出一行对应的注意力结果,即该区间合并后状态的 svC/wdC,结果保留小数点后 6 位。
输入
3 3
0.0 2.0 4.0
0.0 3.0 9.0
0.0 5.0 10.0
2 1 3
2 2 2
2 1 2
输出
2.300000
3.000000
2.600000
说明
三个块的 mx 都为 0,因此任意合并后的 mxC 仍为 0,指数因子 emxi−mxC=e0=1,合并退化为普通的权重和与加权和相加。
查询 2 1 3 时,wd=2+3+5=10,sv=4+9+10=23,注意力结果为 23/10=2.3,输出 2.300000。
查询 2 2 2 只包含第 2 个块,直接计算 sv/wd=9/3=3,输出 3.000000。
查询 2 1 2 时,wd=2+3=5,sv=4+9=13,结果为 13/5=2.6,输出 2.600000。
输入
2 4
1.0 2.0 5.0
0.0 3.0 9.0
2 1 2
1 2 -1.0 1.0 6.0
2 1 2
2 2 2
输出
2.677798
2.721826
6.000000
说明
首次查询 2 1 2 时,合并第 1 块和第 2 块得到 mxC=max(1,0)=1。于是 wd=2e1−1+3e0−1=2+3e−1≈3.103638,sv=5e1−1+9e0−1=5+9e−1≈8.310915,注意力为 sv/wd≈2.677798。
更新操作 1 2 -1.0 1.0 6.0 将第 2 块的状态替换为 (−1,1,6)。再次查询 2 1 2 时,mxC=max(1,−1)=1,wd=2e0+1e−2=2+e−2≈2.135335,sv=5e0+6e−2≈5.812012,注意力为 sv/wd≈2.721826。
查询 2 2 2 为单块查询,最新第 2 块的 wd=1、sv=6,结果为 6/1=6,输出 6.000000。
输入
1 1
-50.0 100.0 -10000.0
2 1 1
输出
-100.000000
说明
该样例只有 1 个块,且查询区间 [1,1] 长度为 1,不涉及块合并。注意力结果为 sv/wd=−10000/100=−100,保留六位小数为 -100.000000。
© CodeFun2000 · 使用条款
Scan the QR code below with WeChat to sign in
First-time scan will create your account automatically
请使用微信扫描下方二维码完成注册