문제

in Blog_old on Boj, Python, Trie

링크

문자열의 집합 두 개 $S_C, S_N$가 주어진다. 각각 크기는 $C$, $N$이다. $Q$개의 문자열 쿼리를 받아서 그 문자열 $q$를 두 부분문자열 $c\in S_C$와 $n\in S_N$으로 분할 할 수 있는지 출력하는 문제이다.

접근

시간 제한 3초에 메모리 1024MB로 제한이 아주 넉넉하다. 백준 온라인 저지는 Python의 경우 여기에다가 추가로 메모리와 시간 제한을 늘려주기 때문에 11초 안에만 풀면 된다🤣

그래서… 예전에 Trie를 알기 전에는 딕셔너리를 사용하여 야매로 풀려고도 했었다.ㅋㅋ

from sys import stdin

c, n = map(int, stdin.readline().split())
colors = {stdin.readline().strip() for _ in range(c)}
names = {stdin.readline().strip() for _ in range(n)}
queries = [stdin.readline().strip() for _ in range(int(stdin.readline().strip()))]
for query in queries:
    flag = False
    for i in range(1, len(query) - 1):
        if query[:i] in colors and query[i:] in names:
            flag = True
            break

    print('Yes' if flag else 'No')

어찌보면 당연하게도 문자열 슬라이싱을 너무 많이 하는 바람에 시간초과를 받았고, 결국 Trie를 공부해서 적용했다.

Trie의 경우 뭐 복잡한 것이 없다. 트리를 만드는데 타고 내려가면서 Prefix를 얻도록 잘 설계하자는 것이 끝이다. 딱 3개의 정보만 관리해주면 된다. 각 노드에 저장할 내용, 그리고 다음 노드가 무엇인지(트리 구조가 되도록), 이 노드에서 끝나는 문자열이 있는지이다.

풀이

나이브하게 딕셔너리를 만들고 잘라보면서 찾자! 는 생각은 시간에 걸렸고, Trie를 적용하기로 했으니 이제 풀이를 구체화해 보자. 쿼리가 들어오면 앞에서부터 한 글자씩 진행하면서 $S_C$로 구성한 Trie를 탐색한다…

그런데 그러면 뒷부분은 어떻게 검색해야 할까? $S_N$으로도 Trie를 구성해서 탐색해야 하나? 생각해보면 Prefix를 그대로 두면서 뒷부분을 붙인 경우를 탐색할때 빠른 것이 이 문제에서 사용하는 Trie의 장점이다. 헌데 뒷부분은 계속 Prefix가 바뀌므로 매번 새롭게 슬라이싱으로 복사해서 사용해야 하고, 그러면 사실상 딕셔너리를 사용하는 것과 다름이 없다.

해서! 나는 그냥 $S_N$에 대해서는 딕셔너리를 사용해버렸다. Best solution을 생각해보자면 아마 문자열을 뒤집어서 Trie를 만들고, 쿼리도 뒤집어서 처리하는 방법으로 두 문자열 집합을 모두 Trie구조로 할 수도 있을 것이다. 그러나 앞서 말했듯 Python은 11초라는 무지막지한 시간 제한이 있기 때문에, 부담없이 복사해버린 다음 딕셔너리에 있는지 확인했다. Trie를 두개 만드는 것보다 이쪽이 간단하다.

Code

링크

# https://www.acmicpc.net/problem/19585
# Trie
from sys import stdin

c, n = map(int, stdin.readline().split())
container = ['']
nxt = [[]]
end = [False]


def add(word):
    pos = i = 0
    while i < len(word):
        for next_pos in nxt[pos]:
            # if next letter matches: move to that node
            if container[next_pos] == word[i]:
                pos = next_pos
                i += 1
                break
        # if nothing matches: insert new node
        else:
            nxt[pos].append(len(container))
            nxt.append([])
            end.append(False)
            pos = len(container)
            container.append(word[i])
            i += 1
    end[pos] = True


# search in trie and yield index whenever available
def find(word):
    pos = i = 0
    while i < len(word) - 1:
        for next_pos in nxt[pos]:
            # if next letter matches: move to that node
            if container[next_pos] == word[i]:
                pos = next_pos
                i += 1
                if end[pos]:
                    yield i
                break
        # if nothing matches: end of iteration
        else:
            return


for _ in range(c):
    color = stdin.readline().strip()
    add(color)

names = {stdin.readline().strip() for _ in range(n)}
for _ in range(int(stdin.readline().strip())):
    query = stdin.readline().strip()
    for i in find(query):
        # print(query[i:])
        if query[i:] in names:
            print('Yes')
            break
    else:
        print('No')

여담

슬슬 코테에 안나오는 알고리즘이라 그런가 솔브닥 class 6 문제임에도 Python으로 푼 사람이 적다. 아마 뒷부분까지 잘 처리해주면 시간을 줄일 수 있을 텐데, 사람이 적어서 제출한 뒤에 1등을 먹어버린터라 귀찮아져버렸다😅😅😅