BEN2のブログ

たまに書いています。

JOI 2015/2016 本選 B - スタンプラリー ２ (Collecting Stamps 2)

競技プログラミング

JOI 2015/2016 本選 B - スタンプラリー ２ (Collecting Stamps 2) についてのメモを残します。

JOI 何個作れる？

まずは、ある文字列が与えられたとき、その文字列から JOI を何個作れるか求めたい。
ある O に注目して、O より左側にある J の個数、O より右側にある I の個数を求めることができれば、注目した O を使って作れる JOI の個数を計算できる。
そこで、下図のように累積和を用意する。

f:id:BEN2suzuka:20200211195414p:plain

下図において、緑色の O に注目すると、それより左側に存在する J の個数が 1 個、右側に存在する I の個数が 2 個であるから、緑色の O で作れる JOI は 2 個、と O(1) で計算できる。

f:id:BEN2suzuka:20200211195424p:plain

ある文字列から作れる JOI の個数を求める計算量は O(N) に抑えられる。

どれをどこに追加する？

追加する文字は J、O、I の 3 通り。それぞれ調べていく。
J は、先頭につけるのが最適。J の累積和の各マスに 1 を加算して、前述の方法で JOI の個数を計算する。
I は、末尾につけるのが最適で、同様に計算できる。

f:id:BEN2suzuka:20200211195436p:plain

O は、文字の間 N-1 箇所のいずれかに配置する。
ある場所に O を追加することで新たに作れる JOI の個数は、下図のように O(1) で求めることができる。
N-1 通りを全て調べ、増える個数の最大値を求めればよい。
文字を追加する前の JOI の個数を別途計算しておく必要がある。

f:id:BEN2suzuka:20200211195445p:plain

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

int main() {
  int N; cin >> N;
  string S; cin >> S;
  vector<long long> J(N+1), I(N+1);
  // J[i] := S[i] より左側に J が何個あるか
  for (int i = 0; i < N; i++) {
    J.at(i+1) = J.at(i);
    if (S.at(i) == 'J') J.at(i+1)++;
  }
  // I[i+1] := S[i] より右側に I が何個あるか
  for (int i = N-1; i >= 0; i--) {
    I.at(i) = I.at(i+1);
    if (S.at(i) == 'I') I.at(i)++;
  }
  long long cnt = 0, c1 = 0, c2 = 0, c3 = 0;
  for (int i = 0; i < N; i++) {
    if (i) {
      // O を入れて増加する JOI 数の最大値
      c1 = max(c1, J.at(i) * I.at(i));
    }
    if (S.at(i) == 'O') {
      // 初期状態の JOI 数
      cnt += J.at(i) * I.at(i+1);
      // 先頭に J をつけたときの JOI 数
      c2 += (J.at(i) + 1) * I.at(i+1);
      // 末尾に I をつけたときの JOI 数
      c3 += (J.at(i)) * (I.at(i+1) + 1);
    }
  }
  long long ans = max(cnt + c1, max(c2, c3));
  cout << ans << endl;
}