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

B - JJOOII 2 (JJOOII 2)

全てのJから条件を満たす文字列を作るのに,どれだけの操作が必要かを考えればよい.それぞれの文字のi番目のindexが分かっていると早く計算できる.'O','I'の開始位置を二分探索で求める. 参考(コピペ)した記事

from bisect import bisect_left
from sys import stdin
input = stdin.readline

n, k = map(int, input().strip().split())
s = input().strip()

index_j = [] #Jのindex
index_o = [] #Oのindex
index_i = [] #Iのindex

for i in range(n):
    if s[i] == 'J':
        index_j.append(i)
    if s[i] == 'O':
        index_o.append(i)
    if s[i] == 'I':
        index_i.append(i)

ans = float('inf')
for start_j in range(len(index_j) - k + 1):
    cnt = index_j[start_j + k - 1] - index_j[start_j] - (k - 1) #start_jからK個のkの間の個数

    start_o = bisect_left(index_o, index_j[start_j + k - 1]) #start_jからK個めのJの位置で二分探索
    if start_o + k > len(index_o): #start_oから後ろにK個の'O'がない
        continue

    cnt += index_o[start_o] - index_j[start_j + k - 1] - 1 #間の文字を数えるので -1
    cnt += index_o[start_o + k - 1] - index_o[start_o] - (k - 1) #start_oからK個のOの間の個数

    start_i = bisect_left(index_i, index_o[start_o + k - 1]) #start_oからK個目のOの位置で二分探索
    if start_i + k > len(index_i): #start_iから後ろにK個の'I'がない
        continue

    cnt += index_i[start_i] - index_o[start_o + k - 1] - 1
    cnt += index_i[start_i + k - 1] - index_i[start_i] - (k - 1)

    ans = min(cnt, ans)

if ans == float('inf'): #条件を満たす文字列が作れない
    print(-1)
else:
    print(ans)

できたこと

  • なし

  • 各文字のindexを前処理で計算しておく

学んだこと

  • 文字列の順番が決まっているときは二分探索で範囲を絞れる
  • 個数をindexの差から求める