[Python] 백준 1893번: 시저 암호

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

1893번: 시저 암호

 

import sys
input = sys.stdin.readline


def getPi(keyword):
    keywordLength = len(keyword)
    pi = [0]*keywordLength
    j = 0

    for i in range(1, keywordLength):
        while j > 0 and keyword[i] != keyword[j]:
            j = pi[j-1]

        if keyword[i] == keyword[j]:
            j += 1
            pi[i] = j

    return pi


def getLeftShiftedChar():
    aLength = len(A)
    leftShiftedChar = {}

    for i in range(aLength):
        leftShiftedChar[A[i]] = A[i-1]

    return leftShiftedChar


def printShifts():
    shifts = []
    aLength = len(A)
    curS = S

    # shifts 추가
    if isThereOneKeyword(curS, W):
        shifts.append(0)

    for i in range(1, aLength):
        curS = shiftLeft(curS)

        if isThereOneKeyword(curS, W):
            shifts.append(i)

    # 출력
    shiftsLength = len(shifts)

    if shiftsLength == 0:
        print("no solution")
    elif shiftsLength == 1:
        print("unique:", shifts[0])
    else:
        print("ambiguous:", end="")
        for n in shifts:
            print("", n, end="")
        print()


def isThereOneKeyword(string, keyword):
    keywordCount = 0
    stringLength = len(string)
    keywordLength = len(keyword)

    i, j = 0, 0

    while i < stringLength:
        if string[i] == keyword[j]:
            i += 1
            j += 1
        else:
            if j == 0:
                i += 1
            else:
                j = pi[j-1]

        if j == keywordLength:
            keywordCount += 1
            if keywordCount >= 2:
                return False

            j = pi[j-1]

    return keywordCount == 1


def shiftLeft(s):
    shifted = s
    length = len(s)

    for i in range(length):
        shifted[i] = leftShiftedChar[shifted[i]]

    return shifted


if __name__ == '__main__':
    N = int(input())

    for _ in range(N):
        A, W, S = list(input().rstrip()), list(
            input().rstrip()), list(input().rstrip())

        pi = getPi(W)
        leftShiftedChar = getLeftShiftedChar()

        printShifts()

 

 

문자열을 shift하는 것은 그리 어렵지 않다. 다만, 문자열에서 검색을 효율적으로 하는 곳이 문제이다. 이를 위해 KMP 알고리즘을 이용하였고, 기타 구현을 생략하고, 이것을 설명하겠다.

 

우선 직관적 이해를 위해 예시를 하나 들겠다.

위의 문자열에서 ABXAB를 찾는다고 하자. 일반적인 문자열 검색은 전체 문자열에 대해 탐색하며, 게다가 ABXAB가 맞는지까지 확인을 해야하므로, 문자열과 검색어의 길이가 N, M이라고 하면, O(NM)이 소요된다.

KMP 알고리즘에선 '패턴'을 이용하여 효율적으로 풀어낸다. 구체적으로는 접두사(prefix)와 접미사(suffix)를 비교하는 것이다.

해당 탐색을 통해, 접미사와 접두사가 AB로 겹친다는 것을 알게 되었다. 이 정보를 이용해, 다음 탐색은 아래와 같이 행한다는 것이다.

 

이제 구체적인 설명을 하겠다.

 

 

우선 키워드에 대해 pi라는 리스트를 설정해주어야 한다. 키워드의 길이가 1~M일 때에 대한 값이다.

즉, pi는 prefix와 suffix가 겹치는 부분의 길이를 저장한다. 이것을 위해 getPi()라는 함수를 작성한다.

def getPi(keyword):
    keywordLength = len(keyword)
    pi = [0]*keywordLength
    j = 0

    for i in range(1, keywordLength):
        while j > 0 and keyword[i] != keyword[j]:
            j = pi[j-1]

        if keyword[i] == keyword[j]:
            j += 1
            pi[i] = j

    return pi

pi[0]은 0이므로 i가 1부터 시작한다.

우선, i, j에 해당하는 값이 같을 경우 j를 증가시키고 pi[i]에 j를 대입한다. 이 경우 for문을 돌게되어 i가 증가하므로, 결국 i와 j를 동시에 1씩 증가시킨 꼴이 된다. 아래 그림처럼 말이다.

 

값이 같지 않을 때는 while문을 돌게 된다. 0이 나오기 전까지, pi[j-1]를 계속해서 넣어주면서 값을 줄여나간다. 여기서 j에 바로 0을 넣지 않는 이유는 다음 2개의 케이스로 설명된다.

화살표는 j = pi[j-1]을 의미한다.

 

이렇게 pi를 완성시켰으면 본격적인 KMP 알고리즘에 의한 검색을 수행한다. 해당 기능은 아래 코드와 같이 수행된다. 다만, 함수의 목적이 키워드가 1개 검색되었는지를 알아내는 것임을 유의하자.

def isThereOneKeyword(string, keyword):
    keywordCount = 0
    stringLength = len(string)
    keywordLength = len(keyword)

    i, j = 0, 0

    while i < stringLength:
        if string[i] == keyword[j]:
            i += 1
            j += 1
        else:
            if j == 0:
                i += 1
            else:
                j = pi[j-1]

        if j == keywordLength:
            keywordCount += 1
            if keywordCount >= 2:
                return False

            j = pi[j-1]

    return keywordCount == 1

i는 string, j는 keyword에 대한 인덱스 값이다.

 

만약 값이 같을 때에는 i, j를 1씩 증가시킨다. 값이 다를 때에는, j에 pi[j-1]을 넣는다. pi를 구할 때와 매우 유사하다.

이렇게 진행되다가, j가 키워드의 길이와 같게 되는 경우가 있는데, 이는 문자열에서 키워드를 하나 찾았다는 것이다. 이 경우에는 keywordCount 관련 동작(함수의 목적을 달성하기 위함으로, KMP와는 별개이다.)을 하고, 아까 한 것과 같이 j에 pi[j-1]을 넣어준다. 아래와 같이 말이다.

즉, j를 prefix와 suffix가 겹치지 않는 곳으로 이동시킨다는 것이다. 이 다음 탐색은 다음과 같이 될 것이다.