[알고리즘] 파이썬으로 Trie 구현하기(2)

 1편에 이어 2편이다. 지난 장에서는 trie의 간단한 개념과 생성과 입력 기능에 대해 알아보았다. 이번 편에서는 조회와 삭제에 대해 알아보도록 하겠다.

접두어로 검색

 Trie의 가장 큰 장점이라고 생각하는 접두어 검색 기능이다. 단어의 일부분으로 전체단어를 검색할 수 있다. 예를들어 trie에 gold, good, gerald를 저장한 다음 검색어에 g를 입력하면 gold, good, gerald를 모두 검색할 수 있다. 이런 기능은 검색어 자동완성에 쓰인다고 한다.
우선 검색 부분의 코드를 보도록 하겠다.

    def search_with_prefix(self, prefix):
        words = list()

        if prefix is None:
            raise ValueError('Requires not-Null prefix')

        top_node = self.head

        for letter in prefix:
            if letter in top_node.children:
                top_node = top_node.children[letter]
            else:
                return words

        if top_node == self.head:
            queue = [node for key, node in top_node.children.items()]
        else:
            queue = [top_node]

        while queue:
            current_node = queue.pop()

            if current_node.data is not None:
                words.append(current_node.data)

            queue = [node for key, node in current_node.children.items()] + queue

        return words

 조회 부분의 원리는 간단하다. 입력한 접두어가 저장되어 있는지 확인하고 해당 접두어가 있다면 그 접두어를 포함하고 있는 모든 단어를 반환해주는 함수이다.
 여기서 queue 부분이 이해가 안갈수도 있을것 같다. 일단 queue 부분은 items()를 이용해 dict의 key, value 값, 즉 노드의 정보를 list에 담아오는 코드이다. 즉, 검색어로 입력한 단어를 포함하는 모든 단어를 찾기위해 검색어로 입력한 이후부분에 해당하는 모든 노드를 찾기 위해 노드 정보를 가져오는 것이다. 그 다음 queue.pop()을 이용해 리스트에 담겨져있는 노드정보를 꺼내 current_node에 저장하고 current_node에 데이터가 있으면 words라는 리스트에 붙이는 과정이라고 이해하면 된다.

삭제

    def delete_word(self, word, depth=0):
        current_node = self.head

        for letter in word[:-1]:
            if letter not in current_node.children:
                return False
            current_node = current_node.children[letter]

        if current_node.NodeCount > 1:
            del current_node.children[word[-1]]
            return True

        if word[-1] in current_node.children:
            del current_node.children[word[-1]]
            self.delete_word(word[:-1], depth)
            return True

        return False

 그 다음은 삭제이다. 삭제는 삭제하고자 하는 단어를 입력하면 그 단어가 있는지 확인하고 공통 접두어가 있는 부분까지 삭제하는 방식이다. 이때 NodeCount가 사용이 되는데 NodeCount가 2 이상이면 자식 노드가 2개 이상이라는 뜻이기 때문에 그 노드에서 삭제를 멈추게하는 방식이다. 하지만 이 삭제 함수의 문제가 하나 있는데 stand와 standard를 구분하지 못한다는 것이다. standard를 삭제하면 stand는 남아있지만, stand를 삭제하면 standard도 삭제한다는 단점이 있다. 이 부분은 차후에 수정을...

if __name__ == '__main__':
    trie = Trie()

    trie.insert_word('good gerald gold')

    print(trie.search_with_prefix('g'))

    trie.delete_word('gold')
    print(trie.search_with_prefix('g'))

이제 작성한 코드를 확인해 볼 시간이다. good gerald gold라는 단어를 trie에 저장하고 g를 이용하여 검색을 해볼 것이다. 그다음 gold라는 단어를 삭제하고 다시 g로 검색을 하는 과정이다.

 출력 결과는 위와 같다. 처음에는 세 단어가 모두 출력되지만 gold를 삭제한 후에는 good과 gerald만 출력된다.