ABC461D Count Subgrid Sum = K 题解

原题链接

Description

有一个 \(H \times W\) 的网格，每个格子包含整数 \(0\) 或 \(1\)。

每个格子中的整数信息由 \(H\) 个长度为 \(W\) 的字符串 \(S_1, S_2, \dots, S_H\) 给出。

如果 \(S_i\) 的第 \(j\) 个字符是 \(0\)，则从上往下第 \(i\) 行、从左往右第 \(j\) 列的格子中写的是 \(0\)；
如果 \(S_i\) 的第 \(j\) 个字符是 1，则该格子中写的是 \(1\)。

请找出满足矩形区域内所写整数的和等于 \(K\) 的矩形区域的数量。

更正式地，找出满足以下所有条件的整数四元组 \((r_1, c_1, r_2, c_2)\) 的个数：

• \(1 \le r_1 \le r_2 \le H\)

• \(1 \le c_1 \le c_2 \le W\)

• 对于满足 \(r_1 \le i ≤ r_2\) 且 \(c_1 \le j \le c_2\) 的所有格子 \((i, j)\) 中所写整数的总和等于 \(K\)。

约束条件：

• \(H\) 和 \(W\) 是 \(1\) 到 \(500\) 之间的整数（包含端点）。

• \(K\) 是 \(0\) 到 \(H\times W\) 之间的整数（包含端点）。

• \(S_i\) 是由 \(0\) 和 \(1\) 组成的长度为 \(W\) 的字符串。

Solution

我们旨在当矩形的顶部和底部行固定时，在 \(\mathcal O(W)\) 时间内解决该问题，从而整个问题可以在 \(\mathcal O(H^2W)\) 时间内解决。

固定顶部和底部行。设 \(A_j\) 为第 \(j\) 列的和，\(B_j = A_1 + A_2 + \cdots + A_j\)。通过适当的预处理，可以在 \(\mathcal O(W)\) 时间内得到这些值。
\(B_r - B_l\) 等于矩形区域中第 \(l+1\) 列到第 \(r\) 列所包含整数的总和，因此我们希望找到满足 \(B_r - B_l = K\) 且 \(0 \le l < r \le W\) 的数对 \((l, r)\) 的个数（其中定义 \(B_0 = 0\)）。

我们应用以下思路来降低复杂度：先统计满足 \(B_r - B_l \ge K\) 的数对个数，再统计满足 \(B_r - B_l > K\) 的数对个数，然后求它们的差。
由于对所有 \(j\) 有 \(A_j \ge 0\)，因此 B 是（弱）单调递增的。对于这两个问题，当 \(l\) 增加时，最小的可行 \(r\) 也（弱）单调递增。我们还注意到，在最小 \(r\) 到 \(W\) 之间的任何整数作为 \(r\) 都是有效的。因此，两个问题都可以用滑动窗口技巧解决。每次滑动窗口扫描的时间复杂度为 \(\mathcal O(W)\)。

总时间复杂度为 \(\mathcal O(H^2W)\)。

Code

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

const int MAX = 500;
int H, W, K, a[MAX + 1][MAX + 1];

i64 solve(int U, int D) {
    vector<int> b(W + 1);
    for (int i = 1; i <= W; ++i)
        b[i] = b[i - 1] + a[D][i] - a[U - 1][i];
    
    i64 cnt = 0;
    for (int l = 1, r1 = 1, r2 = 1; l <= W; ++l) {
        r1 = max(r1, l), r2 = max(r2, l); // *
        while (r1 <= W && b[r1] - b[l - 1] < K)
            ++r1;
        while (r2 <= W && b[r2] - b[l - 1] <= K)
            ++r2;
        cnt += r2 - r1;
    }
    
    return cnt;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    cin >> H >> W >> K;
    
    char ch;
    for (int i = 1; i <= H; ++i) {
        for (int j = 1; j <= W; ++j) {
            cin >> ch;
            a[i][j] = ch - '0';
        }
    }
    
    for (int i = 1; i <= H; ++i)
        for (int j = 1; j <= W; ++j)
            a[i][j] += a[i - 1][j];
    
    i64 ans = 0;
    for (int i = 1; i <= H; ++i)
        for (int j = i; j <= H; ++j)
            ans += solve(i, j);
    
    cout << ans << '\n';
    return 0;
}

\(*\) 处的作用是保证滑动窗口的右端点大于等于左端点。如果不加这行代码，可能会多考虑 \(r<l\) 的非法区间。

\(r<l\) 的情形是可能发生的：如果某次循环初始时 \(l=r\)，而 \(r\) 在循环中没有自增，那么下次循环时就会遇到 \(r<l\) 的非法情形。