とりあえず 1 問目やってみた！累積和の典型題

問題概要

以下のような J, O, I で構成される $H \times W$ の盤面があたえられる。以下の条件を満たすような 3 マスの組が何個あるか求めよ。

3 マスはそれぞれ J, O, I である
J の右側 (行は一緒) に O がある
J の下側 (列は一緒) に I がある

以下の例では 3 組ある。

JOIJ 
JIOO
IIII

制約

$1 \le H, W \le 3000$

考えたこと

'J' が書かれている各マスについて個別に考えていきましょう。各 J マスについて、

その右側にある 'O' の個数を $a$
その下側にある 'I' の個数を $b$

とすると、その 'J' を含む、条件を満たす 3 マスの組の個数は $ab$ 個あります。

ここで問題となるのは、 $a$ や $b$ を各マスごとに毎回愚直に求めていては時間がかかりすぎることです。その時間を見積もってみましょう。

'J' マスの個数： $O(HW)$ 個
その右側や下側にあるマスの探索に要する計算量： $O(H + W)$

ですので、全体の計算量は $O(HW(H+W))$ となります。このままでは TLE となりますので、高速化しましょう。

累積和を用いて前処理

上の解法で時間がかかっているのは、各マスの右側や下側にあるマスの探索です。そこで、次の配列を求めることを考えてみましょう。

onum[i][j]：マス $(i, j)$ の右側にある 'O' マスの個数
inum[i][j]：マス $(i, j)$ の下側にある 'I' マスの個数

これらの配列が求められれば、文字が 'J' である各マス $(i, j)$ について、その 'J' を含む、条件を満たす 3 マスの組の個数は onum[i][j] * inum[i][j] というように、 $O(1)$ の計算量で求められます。

さらに、これらの配列は、実は一回のループで求めることができます。たとえば、onum については、各 $i$ に対して、 $j$ を右側から見ていき、'O' が発生するたびに 1 増やしていく方法で求められます。詳細は、下のコードを参考にしましょう。

計算量

全体の計算量を再度見積もります。

onum, inum を計算するための前処理： $O(HW)$
集計処理： $O(HW)$

となりますので、全体の計算量は $O(HW)$ となります。

コード

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

int main() {
    int H, W;
    cin >> H >> W;
    vector<string> S(H);
    for (int i = 0; i < H; i++) cin >> S[i];

    vector onum(H, vector(W, 0LL)), inum(H, vector(W, 0LL));
    for (int i = 0; i < H; i++) {
        for (int j = W-2; j >= 0; j--) {
            onum[i][j] = onum[i][j+1];
            if (S[i][j+1] == 'O') onum[i][j]++;  // 右側に 'O' が新発生
        }
    }
    for (int j = 0; j < W; j++) {
        for (int i = H-2; i >= 0; i--) {
            inum[i][j] = inum[i+1][j];
            if (S[i+1][j] == 'I') inum[i][j]++;  // 下側に 'I' が新発生
        }
    }

    long long res = 0;
    for (int i = 0; i < H; i++) for (int j = 0; j < W; j++) {
        if (S[i][j] == 'J') {
            res += onum[i][j] * inum[i][j];
        }
    }
    cout << res << endl;
}