Tree 자료구조 심화 - 트라이

1. 문자열 검색

>> 백준 14425번

위 문제를 푸는 방법은 매우 간단하다. 집합 리스트에 대해 in 연산을 수행하여 몇 개의 단어가 집합 S에 포함되어 있는지 출력하면 된다. N과 M의 범위도 10^4까지이기 때문에 이중 for문을 사용해도 시간 제약 2초를 넘기지 않을 것으로 판단된다.

import sys
input = sys.stdin.readline

N, M = map(int, input().split())

src = []
for i in range(N):
    src.append(input().strip())

cnt = 0

for i in range(M):
    word = input().strip()

    if word in src:
        cnt += 1

print(cnt)

하지만, 위와 같은 풀이는 시간 제약이 1초인 경우나, 추가적인 연산이 필요한 문제에서 활용할 수 없다. 따라서, 위 방법을 최적화하는 방법에 대해 알아보기로 하자.

위 코드에서 최적화가 필요한 부분은 in 연산이다. in 연산의 시간 복잡도는 평균적으로 List나 Tuple에서 O(N), Set이나 Dictionary에서 O(1)이다. 참고로, 값을 삽입/삭제할 때의 시간 복잡도 또한 List에서 O(N), Set이나 Dictionary에서 O(1)이다. (물론, List의 끝에서 삽입/삭제하는 경우는 O(1)이다.)

Set과 Dictionary가 List에 비해 in 연산 및 삽입/삭제가 빠른 이유는, Set과 Dictionary가 해시 함수를 통해 원소의 위치를 계산하는 Hash Table 기반 자료구조이기 때문이다. 즉, 위에서 사용한 src 집합을 List 대신 Set이나 Dictionary로 바꿔주면 실행시간이 크게 단축될 것이다. 다만, 딕셔너리를 사용할 경우 value에 무의미한 값을 할당해야 하므로, 여기서는 Set을 활용하기로 한다.

import sys
input = sys.stdin.readline

N, M = map(int, input().split())

src = set()
for i in range(N):
    src.add(input().strip())

cnt = 0

for i in range(M):
    word = input().strip()

    if word in src:
        cnt += 1

print(cnt)

Set 자료구조에 원소를 삽입할 때에는 add() 메서드를 사용한다. 아래의 사진에서 위의 풀이가 Set, 아래의 풀이가 List를 이용한 풀이이다. 보다시피 Set을 사용할 때, 실행 시간이 크게 단축되었다.

2. Trie

1) 개념

사실 위 문제는 Trie 자료구조를 이용하여 풀이할 수도 있다. 본격적인 문제 풀이에 앞서 먼저 Trie 자료구조에 대해 알아보기로 하자.

트라이는 문자열 검색을 빠르게 실행할 수 있도록 설계된 트리 형태의 자료구조로, 주로 문자열 검색 및 자동 완성 기능을 구현하는 데에 사용된다. 트라이는 아래와 같은 특징을 갖는다.

• N진 트리의 형태를 가진다. (여기서 N은 문자의 종류를 의미하는 것으로, 알파벳의 경우 N=26이다.)
• Root 노드는 빈 문자열을 의미하는 공백 상태를 유지한다.

아래는 트라이 자료구조에 apple, air, apply를 차례로 삽입했을 때 나타나는 모습이다.

루트 노드는 공백을 유지한 채로, apple의 각 알파벳에 해당하는 노드를 생성한다. 다음으로 air을 삽입할 때에는 a 노드가 이미 존재하므로 먼저 a 노드로 이동하고, i와 r은 존재하지 않으므로, 새로 생성한다. 마지막으로, apply를 삽입할 때에는 a-p-p-l을 따라 이동하다가, 마지막 y 노드를 새로 생성한다.

2) 구현

① Node 클래스

• 노드에 할당된 문자(self.value), 자식 노드(self.children), 마지막 노드 여부(self.isEnd)를 저장한다.
• 여기서, 마지막 노드 여부는 단어를 정확히 구분하기 위해 필요한 것이다. 예를 들어, Trie에 app과 apple을 모두 삽입했을 때, app이 하나의 독립된 단어로 인정되려면, app도 하나의 단어임을 알리는 Flag가 존재해야 한다.

class Node:

    def __init__(self, value):
        self.value = value
        self.children = {}
        self.isEnd = False

② Trie 클래스

• root 노드가 공백을 유지할 수 있도록, root 노드에는 아무런 값이 없음을 의미하는 Node('')를 할당한다.
• 문자열 삽입 메서드와 문자열 탐색 메서드를 구현한다.

class Trie:

    def __init__(self):
        self.root = Node('') # 공백

    def insert(self, str): 
        node = self.root # 루트에서 시작
        
        for char in str:  # 삽입할 문자열의 각 문자
            if char not in node.children:  # root 노드의 children에 첫 문자가 없으면
                newNode = Node(char)  # 첫 문자를 갖는 새로운 노드 생성
                node.children[char] = newNode  # 생성된 노드를 root 노드의 자식 노드로 추가
                node = newNode  # 자식 노드로 이동
            else:
                node = node.children[char]
        node.isEnd = True  # 단어의 끝 표시

    def search(self, str):
        node = self.root # 루트에서 시작

        for char in str: # 탐색하려는 문자열의 각 문자

            if char in node.children: # root 노드의 children에 첫 문자가 있는 경우
                node = node.children[char] # 자식 노드로 이동

            else:  # root 노드의 children에 첫 문자가 없는 경우
                return False
            
        return node.isEnd # 문자열의 끝에 도달했을 때, 단어도 끝나는지 확인

참고로, 위에서 풀이한 문자열 검색 문제도 Trie 자료구조를 이용하여 해결할 수 있다. (단, Python3로 실행 시 시간 초과가 발생하므로, PyPy3를 사용해야 한다.)

import sys
input = sys.stdin.readline
	
class Node:

    def __init__(self, value):
        self.value = value
        self.children = {}
        self.isEnd = False

class Trie:

    def __init__(self):
        self.root = Node('') # 공백

    def insert(self, str): 
        node = self.root # 루트에서 시작
        
        for char in str:  # 삽입할 문자열의 각 문자
            if char not in node.children:  # root 노드의 children에 첫 문자가 없으면
                newNode = Node(char)  # 첫 문자를 갖는 새로운 노드 생성
                node.children[char] = newNode  # 생성된 노드를 root 노드의 자식 노드로 추가
                node = newNode  # 자식 노드로 이동
            else:
                node = node.children[char]
        node.isEnd = True  # 단어의 끝 표시

    def search(self, str):
        node = self.root # 루트에서 시작

        for char in str: # 탐색하려는 문자열의 각 문자

            if char in node.children: # root 노드의 children에 첫 문자가 있는 경우
                node = node.children[char] # 자식 노드로 이동

            else:  # root 노드의 children에 첫 문자가 없는 경우
                return False
            
        return node.isEnd # 문자열의 끝에 도달했을 때, 단어도 끝나는지 확인
    
N, M = map(int, input().split())
trie = Trie()

for i in range(N):
    element = input()
    trie.insert(element)

cnt = 0

for i in range(M):
    word = input()

    if trie.search(word):
        cnt += 1

print(cnt)

물론, 위 문제는 Trie를 사용할 필요가 전혀 없는 문제이기 때문에, 참고만 하기 바란다.

3) 문제 풀이

이번에는 Trie 구조를 이용해야 하는 문제를 풀어보도록 하자.

>> 백준 5670번

특정 접두사로 시작하는 모든 단어를 찾아야 하는 문제나 문자열을 자동 완성해야 하는 문제는, Trie 자료구조를 사용해 풀이할 수 있다. 먼저, 사용자 입력이 필요한 횟수를 세기 위해, 사용자의 입력이 필요한 경우에 대해 생각해보자.

① 단어의 첫 문자를 입력하는 경우

• 첫 문자는 추론하지 않는다.

② 자식 노드가 유일하지 않은 경우

• hello, hell, heaven이 삽입된 Trie에서 he를 입력한 이후에는 사용자가 a 또는 l을 반드시 입력해야 한다.

③ 어떤 단어의 일부가 독립된 단어가 되는 경우

• hell은 단어일 수도 있고, hello의 일부일 수도 있다.

여기서는 ①과 ③을 한 번에 처리하는 방식을 사용한다. 즉, 단어의 첫 문자가 아닌 마지막 문자에서 사용자의 입력 횟수(cnt)를 증가시킬 것이다. 그 이유는 마지막 문자에서 cnt를 증가시킬 경우, 반드시 사용자가 한번은 입력해야 한다는 조건을 충족하면서, isEnd 노드에 도달했을 때에 cnt를 증가시킬 수 있기 때문이다. (단어의 시작 부분에서 cnt를 증가시킬 경우, isEnd 노드에 도달하는 경우에 대한 처리 로직이 복잡해질 것이다.) 이해를 돕기 위해 hello의 cnt를 세는 과정에 대해 설명하기로 한다.

• 사용자가 처음에 h를 입력하더라도, cnt를 증가시키지 않는다. 이 증분은 단어의 마지막 문자에 도달했을 때 반영할 것이다.
• h의 자식 노드는 유일하므로, 모듈은 e를 자동 완성한다.
• e 노드의 자식 노드는 유일하지 않으므로, 사용자가 l을 입력해야 한다. 이 때, cnt가 1 증가한다.
• l을 입력받으면 l의 자식 노드는 유일하므로, 모듈이 l을 자동 완성한다. 이 때, l은 isEnd가 True인 노드이므로, cnt가 1 증가한다.
• hell이 아닌 hello를 검색하고자 하는 것이므로, 사용자는 o를 입력해야 한다. 이 때, o는 isEnd가 True인 노드이므로, cnt가 1 증가한다.
• 결과적으로 cnt = 3으로 계산된다.

이제 이 내용을 바탕으로 문제를 풀어보자.

import sys
input = sys.stdin.readline
	
class Node:

    def __init__(self, value):
        self.value = value
        self.children = {}
        self.isEnd = False

class Trie:

    def __init__(self):
        self.root = Node('') # 공백

    def insert(self, str): 
        node = self.root # 루트에서 시작
        
        for char in str:  # 삽입할 문자열의 각 문자
            
            if char not in node.children:  # root 노드의 children에 첫 문자가 없으면

                newNode = Node(char)  # 첫 문자를 갖는 새로운 노드 생성
                node.children[char] = newNode  # 생성된 노드를 root 노드의 자식 노드로 추가

            node = node.children[char] # 자식 노드로 이동

        node.isEnd = True  # 단어의 끝 표시

    def count(self, str): # 몇번의 버튼을 눌러야 하는지
        cnt = 0
        node = self.root # 루트에서 시작

        for char in str: # 사용자가 입력하려는 문자열의 각 문자

            if char in node.children:
                node = node.children[char] # 자식 노드로 이동

            if len(node.children) > 1 or node.isEnd: # 자식 노드가 유일하지 않거나, 어떤 단어의 일부가 독립된 단어가 되는 경우 ex) app에선 apple을 바로 완성하지 않아야 함.
                cnt += 1 # 사용자가 버튼을 눌러야 함

        return cnt

while True:

    try: 
        N = int(input())
    except: 
        break

    trie = Trie()
    words = [] # 단어 배열

    for i in range(N):
        word = input().strip()
        words.append(word)
        trie.insert(word)

    cnt = 0

    for word in words:
        cnt += trie.count(word)

    print("%.2f" % (cnt / N)) # "포매팅 코드를 포함한 문자열" % 입력할 값

위 코드에서 아래의 내용을 주의 깊게 살펴보기 바란다.

① try-except

• 입력 값의 개수를 알지 못하는 상황에서 입력이 끝날 때까지, 사용자 입력을 처리하기 위해 사용하는 코드이다.

try:
	N = int(input())
except:
	break

② input().strip()

• 기본적으로 input = sys.stdin.readline는 개행 문자를 포함하기 때문에 입력 값을 그대로 Trie에 삽입하면, 개행 문자가 함께 삽입되는 문제가 발생한다.
• 개행 문자 삽입으로 인한 의도치 않은 동작을 방지하기 위해, Trie 자료구조에 입력 값을 삽입할 때에는 반드시 개행 문자를 제거해주어야 한다.