ABC461 D - Count Subgrid Sum = K
和=Kはいくつ
考え方
愚直にやるのは簡単なコードで実現できる。
つまり、範囲の上下左右の端を決め、その範囲内の総和を求めるループをすればよい。
しかし、それでは $O(H^3W^3)$ かかってTLE。
そこでとりあえず $O(H^2W^2)$ への高速化を目指す。
二次元 vector の矩形範囲内の高速化といえば、二次元累積和。
$O(HW)$ の事前処理をしておけば、任意の範囲の矩形和が $4$ つの数の加算減算だけで求まる。
これで、上下左右の端を決めるループの $O(H^2W^2)$ で処理できるようになる。
しかし、この四重ループで処理するには $157$ 億回のループが必要で、間に合わない。
そこでさらに高速化をする。
範囲の上下だけは同じように決めるとして、$O(H^2)$ 回のループを回す。
このとき、左右を決めるループが、これまでの方法では $O(W^2)$ 回あった。
これを尺取法で $O(W)$ にする。
すなわち、二次元累積和の差を取ることで、その上下幅内の左端から各列までの総和を事前に出しておく。
これは単調増加列になっているはずなので、「値がちょうど $K$ である組がいくつあるか?」は簡単。
尺取り法で「差が $K$ 以下の範囲」「差が $K$ 未満の範囲」を両方持てば、その差が「差が $K$ の範囲」。
ということで、左右を決めるループが $O(W)$ になったので、全体で $O(H^2W)$ となり、計算が間に合う。
入力例1での動作
入力を受け取る。
h: 3
w: 4
k: 3
s: {"1001",
"1101",
"0110"}
まず、グリッド全体の二次元累積和を作る。
{0, 0, 0, 0, 0}
{0, 1, 1, 1, 2}
{0, 2, 3, 3, 5}
{0, 2, 4, 5, 7}
例えば、$2$ 行目から $3$ 行目までを使う場合を考える。
この行の範囲について、左端から各位置までに含まれる 1 の個数を並べると、次のようになる。
{0, 1, 3, 4, 5}
この列方向の累積和を diff とし、右端を $r$ とする。
diff[r]-diff[l] が、その行の範囲で左端 $l$ から右端 $r$ の直前までに含まれる 1 の個数になる。
尺取法では、差が $K$ 以下になる左端の先頭を l1、差が $K$ 未満になる左端の先頭を l2 とする。
すると、l2-l1 が差がちょうど $K$ になる左端の個数になる。
r |
diff[r] |
l1 |
l2 |
この右端での加算 |
|---|---|---|---|---|
| $0$ | $0$ | $0$ | $0$ | $0$ |
| $1$ | $1$ | $0$ | $0$ | $0$ |
| $2$ | $3$ | $0$ | $1$ | $1$ |
| $3$ | $4$ | $1$ | $2$ | $1$ |
| $4$ | $5$ | $2$ | $2$ | $0$ |
この行の範囲では、条件を満たす部分長方形が $2$ 個ある。
他の上下の境界についても同様に処理する。
それぞれで作られる diff と、条件を満たす部分長方形の個数は次のようになる。
| 使用する行 | diff |
条件を満たす個数 |
|---|---|---|
| $1$ 行目 | $\{0,1,1,1,2\}$ | $0$ |
| $1$ 行目から $2$ 行目 | $\{0,2,3,3,5\}$ | $3$ |
| $1$ 行目から $3$ 行目 | $\{0,2,4,5,7\}$ | $2$ |
| $2$ 行目 | $\{0,1,2,2,3\}$ | $1$ |
| $2$ 行目から $3$ 行目 | $\{0,1,3,4,5\}$ | $2$ |
| $3$ 行目 | $\{0,0,1,2,2\}$ | $0$ |
したがって、条件を満たす部分長方形の個数は $0+3+2+1+2+0=8$ となる。
注意点
最大グリッドを全部 $0$ で埋めて $K=0$ とすると、答えは、int 型からはみ出る。
long long 型を用いること。
別解
四重ループで処理するには $157$ 億回のループが必要で、間に合わない。
……と解法に記載したが、実はC++なら間に合わないこともない。
AtCoderのサーバーのCPUは並列処理が可能らしい。
それを活かせるコードを書けば、ギリギリ間に合う。
具体的には、以下をやる。
- 計算は
int型で行う($32$ ビットで処理することで、並列数を上げる) - 連続するループで、
vectorの連続部分を処理していくようにする - ループの中身を単純な計算のみで書く(
.at()は不可)
実際に通った例は上部の「別解」リンク参照。