「3 つのものの真ん中に着目する」という典型が学べる。面白かった！

類題とか

問題概要

"J" と "O" と "I" のみからなる長さ $N$ の文字列 $S$ が与えられます。このような文字列のスコアを、以下の条件を満たす整数 $(i, j, k)$ の組の個数として定めます。

$1 \le i \lt j \lt k \le N$
$S_{i}$ = "J"
$S_{j}$ = "O"
$S_{k}$ = "I"

今、文字列 $S$ の両端および隙間の好きな箇所に "J", "O", "I" のいずれかの文字を挿入して、長さ $N+1$ の文字列にすることができます。そのようにしてできる文字列のスコアの最大値を求めよ。

制約

$3 \le N \le 10^{5}$

考えたこと

まずは「挿入」を無視して、文字列 $S$ のスコアを求める方法を考える。このように 3 つの index $(i, j, k)$ を考察する問題においては、真ん中の添字 $j$ を固定して考えるのが定石。

$S_{j}$ = "O" となるような $j$ を固定したとき、条件を満たす $(i, k)$ の個数は次のように求められる。

$j$ 文字目より左側の範囲の "J" の個数を $a$ 、
$j$ 文字目より右側の範囲の "I" の個数を $b$ としたとき
$a \times b$

$a, b$ は、あらかじめ累積和で管理しておくことで、効率よく求められる。よって、文字列 $S$ のスコアは $O(N)$ で計算できる。

ここから先は、文字列への挿入操作も考慮していく。

"J" を挿入する場合

"J" を追加するときは、"J" を先頭に挿入するのが最適になる。"J" が左にあればあるほど、それより左側の "(O, I)" の個数が多くなるからだ。よって、文字列 "J" + S のスコアを求めれば OK。

"I" を挿入する場合

同様に、"I" を末尾に挿入するのが最適だとわかる。よって、文字列 S + "I" のスコアを求めれば OK。

"O" を挿入する場合

文字列 $S$ を、左右に $l$ 文字 + $N-l$ 文字に分割するように分けてそこに文字 "O" を挿入するとき、スコアは

(左から $l$ 文字の範囲内に含まれる "J" の個数)
× (右から $N-l$ 文字の範囲内に含まれる "I" の個数)

だけ増加する。よって、「文字列 $S$ のスコア」 + 「スコア増加分の最大値」を求めれば OK。

コード

以上をまとめて $O(N)$ でできる。

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

long long calc(string S, bool addO = false) {
    int N = S.size();
    vector<long long> jnum(N+1, 0), inum(N+1, 0);
    for (int i = 0; i < N; ++i) {
        jnum[i+1] = jnum[i] + (S[i] == 'J');
        inum[i+1] = inum[i] + (S[N-1-i] == 'I');
    }
    long long res = 0, addmax = 0;
    for (int i = 0; i < N; ++i) {
        if (S[i] == 'O') res += jnum[i] * inum[N-1-i];
        if (addO) addmax = max(addmax, jnum[i] * inum[N-i]);
    }
    return res + addmax;
}

int main() {
    int N;
    string S;
    cin >> N >> S;
    long long res = max({calc(S, true), calc("J" + S), calc(S + "I")});
    cout << res << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

JOI 本選 2016 B - スタンプラリー２ (AOJ 0626, 難易度 6)

問題概要

制約

考えたこと

"J" を挿入する場合

"I" を挿入する場合

"O" を挿入する場合

コード