JOI 本選 2019 A - 勇者ビ太郎 (AOJ 0658, 難易度 5)

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

問題へのリンク

問題概要

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

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

以下の例では 3 組ある。

JOIJ 
JIOO
IIII

制約

• 

考えたこと

累積和を前処理しておく典型題。各 J に対して、

• 右側に O が何個あるか
• 下側に I が何個あるか

を高速に求めたくて、それは累積和を用いればいい。

コード

#include <iostream>
#include <vector>
#include <string>
using namespace std;

long long onum[3100][3100];
long long inum[3100][3100];

int main() {
    int N, M; cin >> N >> M;
    vector<string> fi(N);
    for (int i = 0; i < N; ++i) cin >> fi[i];

    // O
    for (int i = 0; i < N; ++i) {
        vector<long long> sum(M+1, 0);
        for (int j = 0; j < M; ++j) {
            int add = 0;
            if (fi[i][j] == 'O') ++add;
            sum[j+1] = sum[j] + add;
        }
        for (int j = 0; j < M; ++j) {
            onum[i][j] = sum[M] - sum[j+1];
        }
    }

    // I
    for (int j = 0; j < M; ++j) {
        vector<long long> sum(N+1, 0);
        for (int i = 0; i < N; ++i) {
            int add = 0;
            if (fi[i][j] == 'I') ++add;
            sum[i+1] = sum[i] + add;
        }
        for (int i = 0; i < N; ++i) {
            inum[i][j] = sum[N] - sum[i+1];
        }
    }

    long long res = 0;
    for (int i = 0; i < N; ++i) {
        for (int j = 0; j < M; ++j) {
            if (fi[i][j] != 'J') continue;
            res += onum[i][j] * inum[i][j];
        }
    }
    cout << res << endl;
}