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

H - JOIOJI

まず，各文字の個数で累積和を取る．それぞれの累積和をcsum_j, o ,i とする．区間(i, j)で J, O, I の個数が同じになるとき，csum_j[ j ] - csum_j[ i ] = csum_o[ j ] - csum_o[ i ] = csum_i[ j ] - csum[ i ]となる．
上式を，
csum_j[ j ] - csum_j[ i ] = csum_o[ j ] -csum_o[ i ]，csum_j[ j ] - csum_j[ i ] = csum[ j ] - csum[ i ]
に分けて考えると，
csum_j[ j ] - csum_o[ j ] = csum_j[ i ] - csum_o[ i ]，csum_j[ j ] - csum[ j ] = csum_j[ i ] -csum[ i ]
と変形できる．なので，条件を満たすのは，
csum_j[ j ] - csum_o[ j ] = csum_j[ i ] - csum_o[ i ]，csum_j[ j ] - csum[ j ] = csum_j[ i ] -csum_[ i ]
をともに満たすとき．累積和の差をべつの配列として，同じ値どうしのindexの差の最大値を取ればよい．参考にした記事

from sys import stdin
from collections import defaultdict
input = stdin.readline

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

#各文字の個数の累積和をとる
csum_j = [0]
csum_o = [0]
csum_i = [0]
for i in range(n):
    if s[i] == 'J':
        csum_j.append(csum_j[-1] + 1)
        csum_o.append(csum_o[-1])
        csum_i.append(csum_i[-1])
    if s[i] == 'O':
        csum_j.append(csum_j[-1])
        csum_o.append(csum_o[-1] + 1)
        csum_i.append(csum_i[-1])
    if s[i] == 'I':
        csum_j.append(csum_j[-1])
        csum_o.append(csum_o[-1])
        csum_i.append(csum_i[-1] + 1)

jo = [j - o for j, o in zip(csum_j, csum_o)] #Jの個数=Oの個数
ji = [j - i for j, i in zip(csum_j, csum_i)] #Jの個数=Iの個数

dic = defaultdict(list)
for i, (jo, ji) in enumerate(zip(jo, ji)):
    dic[(jo, ji)].append(i)

ans = 0
for i in dic.values():
    ans = max(i[-1] - i[0], ans)
print(ans)
print(dic)