JOI難易度６ JOIOJI - あっとのTECH LOG

問題原文

問題要旨

J, O, I からなる長さ $N$ の文字列 $S$ が与えられる。
$S$ の連続する部分列であって、その中に含まれる J の数、 O の数、 I の数が等しいようなもののうち、その最大の長さを求めよ。

$1 \leq N \leq 2×10^{5}$

解法

$O(N^{2})$ なら累積和 + 始点終点全探索で終わり。。。なんだけど、制約がきつくてこれは許されない。

とりあえず、J の数、 O の数、 I の数について累積和をとったものを（まぎらわしいけど） $J$ , $O$ , $I$ としておく。
例えば区間[s, t] 内の J と O の数がそれぞれ等しいという条件は、 $J[t\rbrack - J[s - 1\rbrack == O[t\rbrack - O[s - 1\rbrack$ と表現できる。
これを式変形すると、 $J[t\rbrack - O[t\rbrack == J[s - 1\rbrack - O[s - 1\rbrack$ となり、始点と終点を同時に扱わなくても良さそうな形 になる。
これを J と I についても行うと、 $J[t\rbrack - I[t\rbrack == J[s - 1\rbrack - I[s - 1\rbrack$ になる。

この２つが成り立てば与えられる条件を満たせる。（J について消去するように式変形すると、 O と I についても成り立つので）

これは、例えば $(J[t\rbrack - O[t\rbrack, J[t\rbrack - I[t\rbrack)$ をキーとして「どこでこのような状態になるのか」を集計すればいい。そのあとは各キーについて、「最後にそうなったindex」- 「最初にそうなったindex」を計算し、そのMAXが答えになる。

実装

from itertools import accumulate
from collections import defaultdict
N = int(input())
S = input()

J = [0] + list(accumulate([1 if s == 'J' else 0 for s in S]))
O = [0] + list(accumulate([1 if s == 'O' else 0 for s in S]))
I = [0] + list(accumulate([1 if s == 'I' else 0 for s in S]))
JO = [j - o for j, o in zip(J, O)]
JI = [j - i for j, i in zip(J, I)]


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

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

感想

どことなく Zero-Sum-Ranges みがあるなぁと思ったけど自力では解けなかった。
条件式を変形して初めて見えるものがあるということを思い知らされた。。。勉強になりました。。。