基于空间连续块的稀疏注意力机制

题意

给定长度为 $n$ 、维度为 $d$ 的向量序列 $X_0, X_1, \ldots, X_{n-1}$ ，以及块大小 $b$ 和两个 $d$ 维向量 $W_1, W_2$ 。按以下流程计算：

分块：将序列分成 $m = \lceil n/b \rceil$ 个连续块，第 $k$ 块包含 $X_{kb}, X_{kb+1}, \ldots$ （最后一块可能不满）。
块均值：对每个块 $k$ ，计算维度均值 $h_k$ 。
评分：

- $s_k = W_1 \cdot h_k + 2$ （点积加 2）

- $z_k = \max(0,\; s_k)$

- $c_k = W_2 \cdot z_k + \mathbf{1}$ （标量乘法，每个分量再加 1）

- $a_k = \dfrac{\sum_{j} c_{k,j}}{\sqrt{d}}$

二段划分：将序列 $A = (a_1, a_2, \ldots, a_m)$ 分成恰好两段非空连续子段，两段和为 $S_1, S_2$ ，最大化 $\min(S_1, S_2)$ 。

输出 $\text{round}(100 \times S)$ ，其中 $S$ 是最优的 $\min(S_1, S_2)$ 。

思路

这道题本质上就是"按部就班地模拟"，不需要什么高深算法，但有几个小地方值得理一理。

评分公式怎么化简？

先看 $a_k$ 的表达式。 $c_k$ 是一个 $d$ 维向量，第 $j$ 个分量是 $W_{2,j} \cdot z_k + 1$ 。把所有分量加起来：

$ $\sum_j c_{k,j} = z_k \cdot \sum_j W_{2,j} + d$ $

所以 $a_k = \dfrac{z_k \cdot \text{sumW2} + d}{\sqrt{d}}$ ，其中 $\text{sumW2} = \sum_j W_{2,j}$ 只需要预先算一次。这样每个块只需要算一次点积，不用真的构造 $c_k$ 向量。

二段划分怎么做？

把 $A$ 从某个位置切一刀，左边求和得 $S_1$ ，右边求和得 $S_2 = \text{total} - S_1$ 。我们要最大化 $\min(S_1, S_2)$ 。

枚举所有 $m-1$ 个切割点，维护前缀和就行了。当 $S_1$ 从小到大增长时， $S_2$ 从大到小， $\min(S_1, S_2)$ 先增后减，所以取最大值即可。

复杂度

时间： $O(n \cdot d)$ ，计算各块均值和点积
空间： $O(n \cdot d)$ ，存储输入向量

代码

Python3
C++

import sys
import math

def main():
    data = sys.stdin.read().split()
    idx = 0
    n = int(data[idx]); idx += 1
    d = int(data[idx]); idx += 1
    b = int(data[idx]); idx += 1

    X = []
    for i in range(n):
        vec = [float(data[idx + j]) for j in range(d)]
        idx += d
        X.append(vec)

    W1 = [float(data[idx + j]) for j in range(d)]; idx += d
    W2 = [float(data[idx + j]) for j in range(d)]; idx += d

    sumW2 = sum(W2)
    m = (n + b - 1) // b
    sqrt_d = math.sqrt(d)

    A = []
    for k in range(m):
        start = k * b
        end = min(start + b, n)
        bs = end - start
        dot = sum(W1[j] * sum(X[i][j] for i in range(start, end)) / bs for j in range(d))
        s = dot + 2.0
        z = max(0.0, s)
        A.append((z * sumW2 + d) / sqrt_d)

    total = sum(A)
    best = float('-inf')
    prefix = 0.0
    for i in range(m - 1):
        prefix += A[i]
        best = max(best, min(prefix, total - prefix))

    print(round(100 * best))

main()

#include <bits/stdc++.h>
using namespace std;

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, d, b;
    cin >> n >> d >> b;

    vector<vector<double>> X(n, vector<double>(d));
    for(int i = 0; i < n; i++)
        for(int j = 0; j < d; j++)
            cin >> X[i][j];

    vector<double> W1(d), W2(d);
    for(int j = 0; j < d; j++) cin >> W1[j];
    for(int j = 0; j < d; j++) cin >> W2[j];

    double sumW2 = 0;
    for(int j = 0; j < d; j++) sumW2 += W2[j];

    int m = (n + b - 1) / b;
    double sqrt_d = sqrt((double)d);

    vector<double> A(m);
    for(int k = 0; k < m; k++){
        int start = k * b;
        int end = min(start + b, n);
        int bs = end - start;
        double dot = 0;
        for(int j = 0; j < d; j++){
            double col_sum = 0;
            for(int i = start; i < end; i++)
                col_sum += X[i][j];
            dot += W1[j] * col_sum / bs;
        }
        double s = dot + 2.0;
        double z = max(0.0, s);
        A[k] = (z * sumW2 + d) / sqrt_d;
    }

    double total = 0;
    for(int k = 0; k < m; k++) total += A[k];

    double best = -1e18;
    double prefix = 0;
    for(int k = 0; k < m - 1; k++){
        prefix += A[k];
        best = max(best, min(prefix, total - prefix));
    }

    cout << (long long)round(100.0 * best) << endl;
    return 0;
}