【AtCoder】ABC 461 D - Count Subgrid Sum = K

2026-06-08

AtCoder/ABC/D問題 AtCoder/ABC/425点問題 AtCoder/緑Diff AtCoder/NoviSteps/1Q AtCoder/累積和 AtCoder/尺取り法競技プログラミング

D - Count Subgrid Sum = K

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

AtCoder

実行時間制限: 4 sec / メモリ制限: 1024 MiB / Difficulty: 914 / NoviSteps: 1Q / 配点: 425 点

問題概要

$H \times W$ のグリッドがあり、各マスには $0$ または $1$ の整数が書き込まれていて、書き込まれた整数の情報は $H$ 個の長さ $W$ の文字列 $S_1,S_2,\dots,S_H$ として与えられる。

書き込まれた整数の総和が $K$ となる長方形領域がいくつあるか求めよ。

制約

$H,W$ は $1$ 以上 $500$ 以下の整数
$K$ は $0$ 以上 $HW$ 以下の整数
$S_i$ は0, 1からなる長さ $W$ の文字列

考察

いきなり長方形領域の四隅を考えるのは大変なので、まずは上下を固定して考えよう。

上下を $\mathrm{top}$ 行目と $\mathrm{bottom}$ 行目に固定 $(1 \leq \mathrm{top} \leq \mathrm{bottom} \leq H)$ し、 $j$ 列目の値の累積和を $\mathrm{col\_sum}_j$ とする。

すると、求める長方形領域の条件は、 $\displaystyle \sum_{j=l}^{r} \mathrm{col\_sum}_j = K$ となるような区間 $[l, r]$ の個数を求めることに帰着される。

ここで、ある配列 $A$ と整数 $x$ に対して、 $\displaystyle \sum_{i=l}^{r} A_i \geq x$ が成り立つような $(l, r)$ の組の個数を数える関数を $\mathrm{count}(x)$ とすると、求める答えは $\mathrm{count}(K) - \mathrm{count}(K + 1)$ となる。

グリッド上の整数は非負であるため、区間を広げるとその区間内の総和は広義単調増加となる。よって、 $\mathrm{count}(x)$ は尺取り法の要領で計算量 $O(W)$ で求めることができる。

なお、 $K = 0$ の場合は $\mathrm{col\_sum}$ の全ての連続部分列の個数になるので、 $\mathrm{count}(0) = \dfrac{W(W + 1)}{2}$ となる。

また、上下の選び方は二重ループを回せば $O(H^2)$ で全探索できるので、全体の計算量は $O(H^2 W)$ となり、これは本問の制約下で十分高速である。

実装例

1.#include <bits/stdc++.h>
2.using namespace std;
3. 
4.#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)
5. 
6.using ll = long long;
7. 
8.// ======================================== //
9. 
10.int main()
11.{
12.    int H, W, K;
13.    cin >> H >> W >> K;
14.    vector<string> S(H);
15.    rep(i, 0, H) cin >> S[i];
16. 
17.    vector<vector<int>> grid(H, vector<int>(W, -1));
18.    rep(i, 0, H) rep(j, 0, W) grid[i][j] = S[i][j] - '0';
19. 
20.    ll ans = 0;
21.    rep(top, 0, H) {
22.        vector<int> col_sum(W, 0);
23.        rep(bottom, top, H) {
24.            rep(j, 0, W) col_sum[j] += grid[bottom][j];
25. 
26.            auto count = [&](int x) -> ll {
27.                if (x == 0) return W * (W + 1) / 2;
28. 
29.                ll res = 0;
30.                int r = 0, sum = 0;
31.                rep(l, 0, W) {
32.                    while (r < W && sum < x) {
33.                        sum += col_sum[r];
34.                        r++;
35.                    }
36. 
37.                    if (sum >= x) res += W - r + 1;
38.                    sum -= col_sum[l];
39.                }
40. 
41.                return res;
42.                };
43. 
44.            ans += count(K) - count(K + 1);
45.        }
46.    }
47. 
48. 
49.    cout << ans << endl;
50. 
51.    return 0;
52.}