[백준] 26156: 나락도 락이다 (python 파이썬)

우리는 ROCK이라는 단어를 찾고, 그 단어를 포함한 여러 단어에 대한 경우의 수를 구해햐 한다. 이를 위해서는 주어진 문자열을 처음부터 조사하는 것이 아니라 뒤에서부터 살펴봐야 한다. 이때 dp에는 발견한 K, CK, OCK, ROCK의 개수를 dp에 저장하기로 했다. 아래 코드에서 사용한 dp는 입력된 문자열의 길이 n만큼의 길이를 가지며 각 원소는 길이가 5인 리스트로 선언했다. 각 인덱스는 발견한 (0)K, (1)CK, (2)OCK, (3)ROCK의 개수를 저장하고 마지막 인덱스(4)의 칸에는 ROCK을 포함하고 있는 문자열의 개수를 저장했다. 단, ROCK을 포함하고 있는 문자열 중 "ROCK"은 인덱스(3)에 저장되어 있으므로 해당 문자열을 제외한 개수를 저장했다. 즉, "NROCK"이나 "RROCK"과 같이 ROCK에 어떤 문자가 붙어있는 문자의 개수만 포함한다. 이때 다음 문자가 뭔지에 따라 처리할 방법은 아래와 같다.

• 다음 문자가 K인 경우: dp[i][0]=dp[i-1][0]+1
• 다음 문자가 C인 경우: dp[i][1]=dp[i-1][1]+dp[i-1][0] 
• 다음 문자가 O인 경우: dp[i][2]=dp[i-1][2]+dp[i-1][1]
• 다음 문자가 R인 경우: dp[i][3]=dp[i-1][3]+dp[i-1][2]
• 다음 문자가 모든 문자에 대한 처리: dp[i][4] = 2*dp[i-1][4]+dp[i-1][3]

import sys
input = sys.stdin.readline
DIVIDER = 1000000007

n = int(input())
inputString = input().rstrip()

# dp를 생성한다.
dp = [[0 for _ in range(5)] for i in range(n+1)]

# 모든 문자열을 살펴본다.
for i in range(1, n+1):
    dp[i][0] = dp[i-1][0]
    dp[i][1] = dp[i-1][1]
    dp[i][2] = dp[i-1][2]
    dp[i][3] = dp[i-1][3]
    
    # 위에서 설명한 것을 코드로 옮긴다.
    if inputString[n-i]=="K": 
        dp[i][0] = (dp[i-1][0]+1)%DIVIDER
    elif inputString[n-i]=="C":     
        dp[i][1] = (dp[i-1][1]+dp[i-1][0])%DIVIDER
    elif inputString[n-i]=="O": 
        dp[i][2] = (dp[i-1][2]+dp[i-1][1])%DIVIDER
    elif inputString[n-i]=="R": 
        dp[i][3] = (dp[i-1][3]+dp[i-1][2])%DIVIDER
        
    dp[i][4] = (2*dp[i-1][4]+dp[i-1][3])%DIVIDER
    
print((dp[-1][3]+dp[-1][4])%DIVIDER)

위 코드는 dp의 크기가 4×n(문자열의 길이)이다. 하지만 생각해보면 dp의 크기는 단지 4이면 된다. 위 코드는 현재 탐색하는 dp의 위치에 이전 값들을 옮긴다. 즉, 원래 기존의 dp에 이후 과정을 처리해도 아무 문제가 없다. 해당 코드는 아래와 같이 작성할 수 있다.

import sys
input = sys.stdin.readline
DIVIDER = 1000000007

n = int(input())
inputString = input().rstrip()

dp = [0, 0, 0, 0]

for i in range(n-1, -1, -1):
    dp[3] = 2*dp[3]%DIVIDER
    
    if inputString[i] == "K": dp[0] = (dp[0]+1)%DIVIDER
    elif inputString[i] == "C": dp[1] = (dp[1]+dp[0])%DIVIDER
    elif inputString[i] == "O": dp[2] = (dp[2]+dp[1])%DIVIDER
    elif inputString[i] == "R": dp[3] = (dp[3]+dp[2])%DIVIDER
    
print(dp[3]%DIVIDER)