PythonでJOI難易度6を埋める #17

B - スタンプラリー 2 (Collecting Stamps 2)

J は先端に,I は一番後ろに挿入したときに組み合わせは最大になる. O を挿入する場合を考えると,挿入した地点よりも前のJと後ろのIの積だけ増加することが分かります.これは累積和で計算できます.

from sys import stdin
input = stdin.readline

n = int(input().strip())
s = input().strip()

#ここからは先端と一番後ろに J ,I を挿入したときの組み合わせを計算
ans = 0
cnt = 0
cnt_j = [0]
cnt_o = [0]
cnt_i = [0]

s = s + 'I' #実際に結合して数える
for i in reversed(range(n + 1)):
    if s[i] == 'J':
        ans += cnt
        cnt_j.append(ans)
    if s[i] == 'O':
        cnt += cnt_i[-1]
        cnt_o.append(cnt_o[-1] + cnt_i[-1])
    if s[i] == 'I':
        cnt_i.append(cnt_i[-1] + 1)

cnt_j = [0]
cnt_o = [0]
cnt_i = [0]
index_j = []
index_i = []
tmp = 0
cnt = 0

s = s[:-1] #追加した'I'を消去
for i in reversed(range(n)):
    if s[i] == 'J':
        tmp += cnt
        index_j.append(i)
    if s[i] == 'O':
        cnt += cnt_i[-1]
        cnt_o.append(cnt_o[-1] + cnt_i[-1])
    if s[i] == 'I':
        cnt_i.append(cnt_i[-1] + 1)
        index_i.append(i)

ans = max(tmp + cnt_o[-1], ans) #上の二つの場合の大きい方を選択

#ここからは'O'を挿入する場所を考える
count_j = [0]
count_o = [0]
count_i = [0]

for i in range(n):
    if s[i] == 'J':
        count_j.append(count_j[-1] + 1)
        count_o.append(count_o[-1])
        count_i.append(count_i[-1])
    elif s[i] == 'O':
        count_j.append(count_j[-1])
        count_o.append(count_o[-1] + 1)
        count_i.append(count_i[-1])
    elif s[i] == 'I':
        count_j.append(count_j[-1])
        count_o.append(count_o[-1])
        count_i.append(count_i[-1] + 1)

for i in range(1, n): #'O'を挿入する場所
    ans = max(tmp + count_j[i] * (count_i[-1] - count_i[i]), ans)

print(ans)