【带权前缀和】codeforces 2044 H. Hard Demon Problem

题目

https://codeforces.com/problemset/problem/2044/H

题意

第一行，输入一个正整数 \(T(1 \leq T \leq 1000)\)，代表总共 \(T\) 组测试用例。对于每组测试用例：第一行输入两个正整数 \(n, q(1 \leq n \leq 2000, 1 \leq q \leq 10^6)\)，随后输入一个 \(n \times n\) 的矩阵 \(a(1 \leq a_{i, j} \leq 10^6)\)，最后输入 \(q\) 个询问，每个询问输入四个正整数 \(r_1, c_1, r_2, c_2(1 \leq r_1 \leq r_2 \leq n, 1 \leq c_1 \leq c_2 \leq n)\)，代表我们要截取出左上角为 \((r_1, c_1)\) 和右下角为 \((r_2, c_2)\) 的子矩阵。

对于每个询问，我们要计算出 \(1 \times a[r_1][c_1] + 2 \times a[r_1][c_1 + 1] + ... + k * a[r_2][c_2]\) 的值，其中 \(k\) 为子矩阵的元素个数。

保证：所有 \(n\) 之和小于等于 \(2000\)。

题解

对于带权前缀和，可以先学习灵神的视频：https://www.bilibili.com/video/BV1hQ4y1L7Tk/?vd_source=a7061d90bb9470aa775ae3cdb7d30738

对于算式$$1 \times a[r_1][c_1] + 2 \times a[r_1][c_1 + 1] + ... + k * a[r_2][c_2]$$
可以视为：$$\sum_{i=r_1,j=c_1}^{i=r_2, j=c_2}{[(j - c_1 + 1) + (i - r_1) \times (c_2 - c_1 + 1)] \times a_{i,j}} = \sum_{i=r_1,j=c_1}^{i=r_2, j=c_2}{j \times a_{i,j}} + \sum_{i=r_1,j=c_1}^{i=r_2, j=c_2}{(c_2 - c_1 + 1) \times i \times a_{i,j}} + \sum_{i=r_1,j=c_1}^{i=r_2, j=c_2}{[1 - c_1 - r_1 \times (c_2 - c_1 + 1)] \times a_{i,j}}$$

对于 \(\sum_{i=r_1,j=c_1}^{i=r_2, j=c_2}{j \times a_{i,j}}\) 和 \(\sum_{i=r_1,j=c_1}^{i=r_2, j=c_2}{(c_2 - c_1 + 1) \times i \times a_{i,j}}\) 都可以用带权前缀和数组的二维前缀和进行维护；对于 \(\sum_{i=r_1,j=c_1}^{i=r_2, j=c_2}{[1 - c_1 - r_1 \times (c_2 - c_1 + 1)] \times a_{i,j}}\) 只需要维护原矩阵的二维前缀和即可。

时间复杂度：\(O(n^2 + q)\)
空间复杂度：\(O(n ^ 2)\)

参考代码

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

constexpr int N = 2007;
int T, n, q, r1, r2, c1, c2;
ll a[N][N];//前缀和
ll b[N][N];//带权前缀和（列）
ll c[N][N];//带权前缀和的前缀和（b）
ll d[N][N];//带权前缀和（行）
ll e[N][N];//带权前缀和的前缀和（d）

void solve() {
    cin >> n >> q;
    for (int i = 1; i <= n; ++ i) {
        for (int j = 1; j <= n; ++ j) {
            cin >> a[i][j];
            b[i][j] = j * a[i][j];
            b[i][j] += b[i][j - 1];
            c[i][j] = b[i][j] + c[i - 1][j];
            d[i][j] = i * a[i][j];
            d[i][j] += d[i - 1][j];
            e[i][j] = d[i][j] + e[i][j - 1];
            a[i][j] += a[i][j - 1] + a[i - 1][j] - a[i - 1][j - 1];
        }
    }
    while (q --) {
        cin >> r1 >> c1 >> r2 >> c2;
        ll csz = c2 - c1 + 1;
        ll ans = c[r2][c2] - c[r2][c1 - 1] - c[r1 - 1][c2] + c[r1 - 1][c1 - 1];
        ans += csz * (e[r2][c2] - e[r2][c1 - 1] - e[r1 - 1][c2] + e[r1 - 1][c1 - 1]);
        ans += (1LL - r1 * csz - c1) * (a[r2][c2] - a[r2][c1 - 1] - a[r1 - 1][c2] + a[r1 - 1][c1 - 1]);
        cout << ans << ' ';
    }
    cout << '\n';
}

int main() {
    ios::sync_with_stdio(false);cin.tie(nullptr);cout.tie(nullptr);
    cin >> T;
    while (T --) {
        solve();
    }
    return 0;
}