백준 19585번: 전설 (Trie)

문제링크:19585번: 전설 (acmicpc.net)

19585번: 전설

이 문제는 트라이를 이용하여 문자열을 탐색해 나가는 문제이다. 특이한 점으로는 color을 뜻하는 문자열 C개, name을 뜻하는 문자열 N개가 주어지고 팀 명이 주어지는데, 팀명이 반드시 color 문자열 + name 문자열 형식으로 이루어 져 있는지 판별해야 한다. C, N은 모두 최대 4000개까지 주어질 수 있고, color, name의 문자열 개수는 최대 1000자 이다. 팀의 개수 Q는 20000개까지 주어지고, 팀의 문자열 개수는 최대 2000자이다.

 

만약 브루트 포스로 탐색 한다면 시간 복잡도가 20000 * 4000 * 4000으로 무조건 시간 초과가 날 것이므로 트라이를 사용했다. 먼저 color 문자열들을 ctrie에 트라이 형태로 저장하고 name 문자열들을 ntrie에 트라이 형태로 저장한다. 여기서 ntrie는 문자열의 순서를 역순으로 탐색해 가며 트라이로 저장한다. 그 이유는 ntrie를 순차적으로 탐색한다면, 팀명을 탐색할 때에 앞에서부터 color을 탐색하고, color이 완성되면 완성될 때마다 ntrie를 조사해야 한다.

 

예시로

color : 'red' 'redo' 'redoo' 라고 있을때, 팀명이 color에서 'red', 'redo', 'redoo' 를 사용한 지 모르기 때문에 3가지 경우에서 모두 ntrie 탐색을 시작해야 한다.

 

이 문제를 해결하기 위해 ntrie를 역순으로 저장, ntrie를 팀명의 역순부터 시작해 탐색해 나간다면 ntrie 탐색을 한번만 하고, 가능한 name들의 인덱스들을 저장하고, ctrie를 한번 탐색하며 가능한 color들의 인덱스를 저장하면 가능한 color의 인덱스, name의 인덱스들을 비교해 가면 ctrie, ntrie탐색 모두 한번만 해서 답을 구할 수 있다.

 

처음에 나는 전에 사용했던 재귀를 이용해 트라이를 탐색하여 문제를 풀어 보았는데, 시간 초과가 나왔다. 그래서 재귀함수를 for문으로 바꾸어 트라이 탐색해 보았는데, 또 시간 초과가 나왔다.

 

시간 초과 풀이:

import sys

C,N=map(int,sys.stdin.readline().rstrip().split())

ctrie={}
ntrie={}

for _ in range(C):
    color=list(sys.stdin.readline().rstrip())
    cur=ctrie
    for letter in color:
        if letter not in cur:
            cur[letter]={}
        cur=cur[letter]
    cur['*']=1

for _ in range(N):
    name=list(sys.stdin.readline().rstrip())
    name=reversed(name)
    cur=ntrie
    for letter in name:
        if letter not in cur:
            cur[letter]={}
        cur=cur[letter]
    cur['*']=1

Q=int(sys.stdin.readline().rstrip())

for _ in range(Q):
    findans=0
    team=list(sys.stdin.readline().rstrip())
    k = len(team)
    colorends=[0 for _ in range(2000)]
    
    cur=ctrie
    for idx,letter in enumerate(team):
    	if idx >= 1000:
        	break
        if letter not in cur:
            break
        cur=cur[letter]
        if '*' in cur:
            colorends[idx]=1

    cur=ntrie
    for idx,letter in enumerate(reversed(team)):
    	if idx >= 1000:
        	break
        if letter not in cur:
            break
        cur=cur[letter]
        if '*' in cur:
            if colorends[k-2-idx]:
                print('Yes')
                findans=1
                break
                
    if findans==0:
        print('No')

시간 복잡도는 O(Q * (C+N)) 로 40,000,000 이다. 40,000,000이면 확실히 통과를 보장할 순 없는 정도라 최적화를 더 잘 해야 될 것이라 생각해서 파이썬 dictionary의 탐색 과정에 대해 조사해 보았다.

 

TimeComplexity - Python Wiki

TimeComplexity - Python Wiki

파이썬 dictionary는 각 키들을 hash형태로 저장하여 탐색한다. 만약 hash내부의 충돌이 일어나면, a in dict 함수는 시간 복잡도가 O(1)이 아닐 수도 있다고 한다. 근데 우리는 키값으로 'a', 'b', 'c'.. 'z' 등 알파벳 소문자만 사용해서 hash충돌이 웬만하면 일어나지 않을 것이다. 실제로 알파벳 소문자들의 hash값을 확인 해 보면 아래와 같게 나온다.

abc = list('abcdefghijklmnopqrstuvwxyz')
for letter in abc:
    print(hash(letter))

***
-8303460262433247325
-4112022886574324644
-401239679903848455
-2045380628259562936
-3875376347303506065
-9164622502464746828
-3660767190487437868
-2483359002667067152
-3547501648993601227
-6787663625330551333
8084369612867151960
-5714842814981535235
-455803042856853688
5381949143570322950
2795783468299795288
3371996166798448844
-2370857241689283464
-2098513289873047545
-6582113614864899633
5091762767117415136
-147173317997526726
-4363236706252780946
-2867165210419187236
-219003678142575924
4170892727866127602
8062143613808912215
***

그럼 해시 충돌은 아닌데, 위의 hash값들을 보고 최적화 할 방법을 찾았다. 위와 같이 hash값은 int64비트를 사용하는 것을 알 수 있다. 그런데 우리가 실제로 사용할 문자는 알파벳 26자 + '*' 이므로 27개이므로 굳이 int64비트를 사용하여 저장해 메모리와 시간을 늘릴 필요가 없다. 그래서 letter마다 ord(letter) - 97 해주어 'a', 'b', 'c' ..'z' 를 0, 1, 2, 3... 26 로 변환하여, 문자열의 끝을 알리는 '*' 대신 정수 27을 대입하여 ctrie, ntrie에 저장하였다. 이렇게 하면 각 정수의 해시값은 그 정수 (hash(1) = 1, hash(2) = 2..) 이기때문에 매번 트라이를 탐색할 때마다 시간을 줄일 수 있다.

 

import sys

C,N=map(int,sys.stdin.readline().rstrip().split())


ctrie={}
ntrie={}

for _ in range(C):
    color=list(sys.stdin.readline().rstrip())
    cur=ctrie
    for letter in color:
        letter = ord(letter)-97
        if letter not in cur:
            cur[letter]={}
        cur=cur[letter]
    cur[27]=1
        
for _ in range(N):
    name=list(sys.stdin.readline().rstrip())
    name=reversed(name)
    cur=ntrie
    for letter in name:
        letter = ord(letter)-97
        if letter not in cur:
            cur[letter]={}
        cur=cur[letter]
    cur[27]=1

Q=int(sys.stdin.readline().rstrip())

for _ in range(Q):
    findans=0
    team=list(sys.stdin.readline().rstrip())
    k = len(team)
    colorends=[0 for _ in range(2000)]
    
    cur=ctrie
    for idx,letter in enumerate(team):
        if idx >= 1000:
            break
        letter = ord(letter)-97
        if letter not in cur:
            break
        cur=cur[letter]
        if 27 in cur:
            colorends[idx]=1

    cur=ntrie
    for idx,letter in enumerate(reversed(team)):
        if idx >= 1000:
            break
        letter = ord(letter)-97
        if letter not in cur:
            break
        cur=cur[letter]
        if 27 in cur:
            if colorends[k-2-idx]:
                print('Yes')
                findans=1
                break
                
    if findans==0:
        print('No')

이렇게 헀더니 맞았습니다를 받을 수 있었다. 파이썬의 해시에 대해 알 수 있어서 좋은 문제였다. 만약 dictionary나 set을 사용할 일이 있을때, 시간을 줄이고 싶을 때 여기에 저장될 값들의 개수를 알 수 있다면 이렇게 정수로 변환하여 저장하는 것도 좋은 것 같다. 만약 알파벳 대문자까지 이용한다면 0 ~ 52의 값을 대입 하여 할 수 있다. 또는 간단하게 그냥 아스키 코드를 사용해도 될 것이다.