Trie 자료구조와 문제풀이 [leetcode]

Leetcode 에서 매일 하나씩 문제가 나오고 그걸 풀다보면 매주 마다 하나씩 테마가 있다고 느끼게 된다.

 

이번주는 trie라는 자료구조를 써야하는 문제가 자주 나왔고, 마침 내가 몰랐던 자료구조이기에 이번 기회에 한번 정리하고 넘어가려 한다. (물론 이전에도 한두문제 나왔었지만 그냥 포기하고 넘어갔었다.)

 

위키피디아를 긁어오자면,

? 머리에 물음표가 뜬다. 무슨소리인가. 트리인건 알겠는데.

같이 그려져 있는 그림을 가져와보자

그림을 살펴보면 ‘t’라는 단어로 시작되는 단어가 있는지(exist), 있다면 뭐가 있는지(search), 몇개 있는지(how many) 등을 트리를 통해 빠르게 접근가능하다는 장점이 있는 자료구조이다.

어떤 블로그를 보니 Retrieval 에서 trie라는 단어를 따왔다던데 단어 검색에 장점이 있는 자료구조가 되겠다.

 

1. Implement Trie (Prefix Tree)

trie 자료구조에 Insert를 구현하고 해당 자료구조를 사용해서 search와 startWith를 구현하는 문제다.

우선 trie에서 연결이 될 각 Node를 정의한다.

class Node:
    def __init__(self):
        self.data = [] # 각 노드에 담기는 문자열을 저장
        self.children = {} # 하위의 시작 문자열을 저장

각 노드를 통해 저장하고자 하는 모습을 먼저 상상해보자.

“apple” 이라는 단어가 들어왔을때

root

• a : data = []
  • p : data = []
    • p : data = []
      • L : data = []
        • e : data = [”apple”]

“app” 이라는 단어가 들어왔을때,

root

• a : data = []
  • p : data = []
    • p : data = [”app”]
      • L : data = []
        • e : data = [”apple”]

각 하위 값들로 다음에 오는 알파벳을 저장하면 startWith를 조회할때 a, ap, app, appl 등으로 시작하는 단어가 존재하는지 여부를 알수 있게 되고(children), 각 Node의 data 값을 통해 단어 존재여부를 파악할 수 있게 된다.

완성된 코드로 살펴보자.

class Trie:

    def __init__(self):
        self.head = Node()

    def insert(self, word: str) -> None:
        current_node = self.head
        for w in word: # 각 자리의 알파벳별로 순회
            if w not in current_node.children: # 해당 알파벳의 현재 노드의 자식으로 없다면
                current_node.children[w] = Node() # 새로운 Node 생성해준뒤,
            current_node = current_node.children[w] # 자식 Node로 이동
        current_node.data.append(word) # 마지막 자식노드에 단어를 data로 삽입
        
    def search(self, word: str) -> bool:
        current_node = self.head
        for w in word: # 각 자리의 알파벳별로 순회
            if w not in current_node.children: # 만약 중간의 자식노드가 없다면
                return False # 존재하지 않는 단어 이므로 False 반환
            else:
                current_node = current_node.children[w]
        if word in current_node.data:
            return True
        else:
            return False

    def startsWith(self, prefix: str) -> bool:
        current_node = self.head
        for w in prefix: #prefix의 각 자리 알파벳별로 순회
            if w not in current_node.children: # 중간의 자식노드가 없다면
                return False # 존재하지 않는 접두사이므로 False 반환
            else:
                current_node = current_node.children[w]
        # 만약 자식노드가 존재하거나, 현재 노드의 단어데이터가 존재한다면 접두사가 맞으므로 True
        if len(current_node.children) > 0 or len(current_node.data) > 0:
            return True
        else:
            return False

class Node:
    def __init__(self):
        self.data = []
        self.children = {}
# Your Trie object will be instantiated and called as such:
# obj = Trie()
# obj.insert(word)
# param_2 = obj.search(word)
# param_3 = obj.startsWith(prefix)

전체 코드와 문제는 여기서 볼 수 있다.

https://github.com/hoonzinope/algorithm_study/tree/main/leetcode/Implement-Trie-(Prefix%20Tree) 

GitHub - hoonzinope/algorithm_study: algorithm 문제풀이

 

다른 문제도 풀어보자

 

2. Prefix and Suffix Search

마찬가지로 단어 검색문제다.

 

들어온 단어 중에 검색하려는 prefix와 suffix(위 문제와 다른부분) 을 검색하는 건데,

위 문제와 동일하게 prefix로 Node 를 타고 내려온 다음

suffix가 등장하는 단어를 찾는 문제가 되겠다. (그리고 해당 단어의 index를 반환해야한다.)

 

이번엔 좀 다르게 풀었다. 예시를 들어보자면

“apple” 이라는 단어가 들어왔을 때,

root

• a : data = [”apple”]  #prefix 로 등장하는 단어들을 전부 적어줌
  • p : data = [”apple”]
    • p : data = [”apple”]
      • L : data = [”apple”]
        • e : data = [”apple”]

“app” 이라는 단어가 들어왔을 때,

root

• a : data = [”apple”,”app”]
  • p : data = [”apple”,”app”]
    • p : data = [”apple”,”app”]
      • L : data = [”apple”]
        • e : data = [”apple”]

prefix에 해당되는 단어를 전부 저장해서 내려가는 방법으로 풀었다. 구조에 이렇게 저장하게 되면

“ap” 라는 prefix를 검색시 Node를 타고 내려온 뒤엔 해당 Node에 data 부분에서 검색이 가능하다는 장점이 있다.

(dfs같이 Node 를 타고 내려가는 대신에!)

 

코드를 보자.

# 1번 문제보다 이전에 풀어서 다른 블로그를 보고 풀었더니, Node에 필요없는 key값을 넣어버림
class Node: 
    def __init__(self, key, data=None):
        self.key = key
        self.data = []
        self.children = {}

위 1번 문제와 마찬가지로 Node 값을 설정해준다.

class WordFilter:
    def __init__(self, words: List[str]):
        self.head = Node(None)
        self.result_dict = {} # 똑같은 search 요청이 올경우 해당 결과값을 저장해둘 dictionary
        
        for index, word in enumerate(words):
            current_node = self.head
            current_node.data.append((word, index)) # 맨처음 root Node에는 모든 단어를 저장
            for alpha in word:
                if alpha not in current_node.children: 
                    current_node.children[alpha] = Node(alpha)
                # 위 문제와 다른점은 각 Node에 단어들을 저장, prefix로 접근시 단어를 data에서 조회 가능
                current_node.children[alpha].data.append((word, index))
                current_node = current_node.children[alpha]

 

self.result_dict 라는 dictionary 변수를 하나 뒀는데,

문제에서 동일한 검색으로 엄청 검색해서 들어오는 케이스가 하나 있어서

한번 검색된 결과값을 저장해두기 위해 생성했다.

 

위 1번 문제랑 다른 점은 각 Node에 타고 들어갈 때 마다 data에 해당 단어들을 전부 저장해준다는 점이다.

prefix를 통해 들어온 Node 내 data 리스트에서 suffix를 검색하게끔 만들기 위함이다.

 

WordFilter 클래스의 검색부분 함수를 마저 살펴보자.

def f(self, prefix: str, suffix: str) -> int:
    current_node = self.head

    if (prefix,suffix) in self.result_dict: # 만약 검색결과가 존재할 경우, 해당 결과값을 반환
        return self.result_dict[(prefix,suffix)]
    
    for pre in prefix: # prefix를 통해 Node를 타고 넘어가기
        if pre in current_node.children:
            current_node = current_node.children[pre]
        else: # 중간에 타고갈 Node가 없다면 존재X
            return -1
    
    max_num = -1 # 존재한다면 data 중에 index가 가장 큰걸 반환해달라고 했으니 
    for word, index in current_node.data: # data를 순회하면서 조건에 맞는 부분을 반환
        if suffix == word[-len(suffix):]:
            if max_num < index:
                max_num = index
    
    self.result_dict[(prefix,suffix)] = max_num # 검색기록을 저장 (나중에 또물어보면 여기서 꺼내씀)
    return max_num

위에 init 부분에서 정의한 self.result_dict 에

현재 검색하려는 부분이 존재한다면 반환한다.

(물론 이건 중간에 추가되는 단어가 없기때문에 가능한 일이다.)

 

검색되는 결과가 없다면 prefix를 알파벳 단위로 순회하면서 Node를 타고 내려간다.

중간에 타고 내려갈 Node가 없다면 prefix에 맞는 단어가 없다는 뜻이므로 -1 을 반환

 

존재한다면, Node의 data를 돌면서 suffix조건에 맞는 단어를 조회해준다.

(문제에서는 가장 큰 index를 반환하라고 했으므로 max_index를 찾아준다.)

 

전체 코드와 문제는 여기서 볼 수 있다.

https://github.com/hoonzinope/algorithm_study/tree/main/leetcode/Prefix-and-Suffix-Search

GitHub - hoonzinope/algorithm_study: algorithm 문제풀이

 

마지막 문제~ 난이도 Hard

3. Search Suggestions System

난이도 hard의 문제였다.

trie 문제를 이렇게 쭉 풀어왔으면,

‘아하 trie쓰면 금방 풀겠네’ 했겠지만 그게 아니였으면 어떻게 풀어야 할지 감이 안잡히는 문제여서 Hard 인가보다.

 

마찬가지로 Node를 만들고

단어가 들어올 때마다 각 자리 알파벳 별로 순회하면서 Node Tree를 만들고

각 Node별로 data 리스트를 둬서 나오는 값들 반환하면 된다.

 

코드를 보자

class Node:
    def __init__(self, key, data=None):
        self.key = key
        self.data = []
        self.children = {}

class Solution:
    def __init__(self):
        self.head = Node(None)
    
    def suggestedProducts(self, products: List[str], searchWord: str) -> List[List[str]]:
        products = sorted(products)
        
        for product in products:
            current_node = self.head
            for alpha in product:
                if alpha not in current_node.children:
                    current_node.children[alpha] = Node(alpha)
                    current_node.children[alpha].data.append(product)
                else:
                    current_node.children[alpha].data.append(product)
                current_node = current_node.children[alpha]
        
        result = []
        current_node = self.head
        for alpha in searchWord:
            if alpha in current_node.children:
                if len(current_node.children[alpha].data) > 3:
                    result.append(current_node.children[alpha].data[:3])
                else:
                    result.append(current_node.children[alpha].data)
                current_node = current_node.children[alpha]
            else:
                result.append([])
                current_node.children[alpha] = Node(None)
                current_node = current_node.children[alpha]
        return result

 

products로 들어오는 문자값들을 하나씩 돌면서 Node를 생성해준다.

Node 가 생성될때마다 data 리스트에는 해당 문자열 전체(product)가 저장된다.

 

suggestedProducts 함수를 보면 문제에서 요구하는 바를 만족시키기 위해

searchWord를 알파벳 별로 돌면서 해당 값을 결과로 저장한다.

“mouse”를 검색한다고 할때, m → o → u → s → e 순으로 Node를 순회하고,

각 Node에 도달시 마다 저장된 data 리스트를 결과에 저장해준다.

 

전체 코드와 문제는 여기서 볼 수 있다.

https://github.com/hoonzinope/algorithm_study/tree/main/leetcode/Search-Suggestions-System

GitHub - hoonzinope/algorithm_study: algorithm 문제풀이