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は並列処理が可能らしい。
それを活かせるコードを書けば、ギリギリ間に合う。

具体的には、以下をやる。

実際に通った例は上部の「別解」リンク参照。