[Python] 5525번 IOIOI

https://www.acmicpc.net/problem/5525

5525번: IOIOI

---

 

22/09/14

 

 

은근 아이디어를 요구했던 문제이다. 창의적이고 재미있었던 문제였다.

---

 

문제 접근 방식:

 

 

처음에는 문자열 S안에 Pn이 몇 개 있는지 물어봐서 인덱스를 하나씩 돌려가며 Pn이랑 매칭이 되는지 판별하는 알고리즘을 짰었다.

 

그랬더니 문자열의 길이가 M이므로 대략 O(MN) 만큼의 시간 복잡도가 걸린다는 것을 유추할 수 있었다.

 

하지만 시간초과가 걸렸고.... 나는 다른 알고리즘을 짜야만 했다.

---

아래는 내가 위의 접근 방식과 같이 작성한 파이썬 코드이다. 시간 초과가 난 코드이지만, 좋은 참고가 될 수 있을 것 같아 올려본다. 더보기를 누르면 확인할 수 있다.

# 5525번 IOIOI
# 문자열
import sys
input = sys.stdin.readline

N = int(input())
M = int(input())
S = input().rstrip()
total = 0
for start in range(M-N):
    for i in range(2*N+1):  # Pn의 길이 만큼 반복하여 돈다.
        if i % 2 == 0:  # Pn의 글자가 I
            if S[start+i] != 'I':
                total -= 1
                break
        else:
            if S[start+i] != 'O':
                total -= 1
                break
    total += 1
print(total)

 

---

나는 Pn이라는 문자열이 I와 O로만 이루어져있고 굉장히 규칙적이라는 점에 대해 주목을 했다.

 

그러다 한가지 사실을 발견했는데, 만약 Pk이 주어져 있고, Pn가 주어져 있으면(k >= n) Pk안에 Pn가 몇 개 들어갈 수 있나?를 따져 보면 k - n + 1번 들어갈 수 있음을 알아냈다.

 

그렇다면, 우리는 S라는 문자열을 인덱스를 늘려가며 문자 하나하나를 확인하는데, IOIOI가 끊기는 그 순간, 그때까지의 문자열이 Pk이라고 할 것이다.

 

만약 Pk가 Pn보다 길다면, 위의 공식을 통해 몇 번 들어가는지를 O(1)의 시간 복잡도로 구할 수 있으므로, 우리는 O(MN)이 아니라 O(M)의 시간 복잡도로 구할 수 있을 것이다.

 

예를 들어, OOIOIOIOIIOII 라고 한다면, 2번 인덱스부터 IOIOI문자열이 시작이 되는 것이다.

 

IOIOI문자열이 어디서 끝나냐면 8번 인덱스에서 끊기는데, 그러면 IOIOI문자열은 IOIOIOI, 즉, P3에 해당이 되는 것이다.

그러면 여기서 k는 3에 해당이 된다. 이 k값을 구하는 방법은 (IOIOI문자열의 길이)//2를 해주면 구할 수 있다.

 

믈론 IOIOI문자열이 위의 상황과는 다르게 I가 겹쳐서 끊기는 것이 아닌, O가 겹쳐서 끊길 수도 있다.

 

예를 들자면, OOIOIOIOIOOII 라고 하면, 2번 인덱스부터 9번 인덱스까지가 IOIOI문자열의 길이이다.

근데, 이는 P3이랑 완전히 같지는 않은데, 이는 나눗셈이 될 때 내림을 하면 되는 부분으로 괜찮다.

---

아래는 내가 위의 접근 방식과 같이 작성한 파이썬 코드이다. 더보기를 누르면 확인할 수 있다.

# 5525번 IOIOI
# 문자열
'''
접근 방식:
Pk이 주어져 있고, Pn가 주어져 있으면(k >= n)
Pk안에 Pn가 몇게 들어갈 수 있나?를 따져 보면
k - n + 1번 들어갈 수 있음을 이용한다.
'''
import sys
input = sys.stdin.readline

N = int(input())
M = int(input())
S = input().rstrip()
total = 0
Pk_len = 0
for i in range(M):
    if Pk_len % 2 == 0 and S[i] == 'I':
        Pk_len += 1
    elif Pk_len % 2 == 1 and S[i] == 'O':
        Pk_len += 1
    else:
        k = (max(0, Pk_len - 1)) // 2
        if k >= N:
            total += k-N+1
        Pk_len = 0
        if Pk_len % 2 == 0 and S[i] == 'I':
            Pk_len += 1
k = (max(0, Pk_len - 1)) // 2
if k >= N:
    total += k-N+1
print(total)