Computer >> 컴퓨터 >  >> 프로그래밍 >> 프로그래밍

파이썬 트라이(Trie) 자료구조로 자동 완성 기능 구현하기

자동 완성 기능이란?

검색창에 글자를 입력할 때마다 관련 단어가 실시간으로 추천되는 경험은 누구나 한 번쯤 겪어봤을 것입니다. 이처럼 사용자가 입력한 문자와 일치하는 문자열을 즉시 보여주는 기능을 자동 완성(Auto Completion)이라고 합니다.

예를 들어 트라이(Trie)에 "xyz", "xyzzzz", "xxxyyxzzz"라는 단어가 저장되어 있고, 사용자가 xy를 입력했다면 화면에는 xyz, xyzzzz 등 접두사가 일치하는 모든 단어를 보여줘야 합니다.

자동 완성 구현 절차

  • 표준 트라이 탐색 알고리즘으로 입력된 문자열을 검색합니다.

  • 문자열이 존재하지 않으면 -1을 반환합니다.

  • 문자열이 존재하고 그 지점이 트라이 내 단어의 끝이라면 해당 문자열을 출력합니다.

  • 일치하는 지점 뒤에 더 이상 하위 노드가 없다면 종료합니다.

  • 그 외의 경우에는 해당 노드 아래에 연결된 모든 노드(완성 가능한 단어들)를 출력합니다.

파이썬 코드 예제

# 트라이 노드 클래스
class TrieNode():
    def __init__(self):
        # 트라이 노드 초기화
        self.trie_node = {}
        self.last_node = False

class Trie():
    def __init__(self):
        # 트라이 초기화
        self.root = TrieNode()
        # 단어를 저장할 리스트
        self.words = []

    def create_trie(self, keys):
        # 주어진 데이터로 트라이 생성
        for key in keys:
            # 트라이에 키 하나씩 삽입
            self.insert_node(key)

    def insert_node(self, key):
        node = self.root
        for obj in list(key):
            if not node.trie_node.get(obj):
                # 새로운 TrieNode 생성
                node.trie_node[obj] = TrieNode()
            node = node.trie_node[obj]
        # 리프 노드 표시
        node.last_node = True

    def search(self, key):
        # 키 검색
        node = self.root
        is_found = True
        for obj in list(key):
            if not node.trie_node.get(obj):
                is_found = False
                break
            node = node.trie_node[obj]
        return node and node.last_node and is_found

    def matches(self, node, word):
        if node.last_node:
            self.words.append(word)
        for obj, n in node.trie_node.items():
            self.matches(n, word + obj)

    def show_auto_completion(self, key):
        node = self.root
        is_found = False
        temp = ''
        for obj in list(key):
            # 단어 존재 여부 확인
            if not node.trie_node.get(obj):
                is_found = True
                break
            temp += obj
            node = node.trie_node[obj]
        if is_found:
            return 0
        elif node.last_node and not node.trie_node:
            return -1
        self.matches(node, temp)
        for string in self.words:
            print(string)
        return 1

# 트라이에 저장할 데이터
strings = ["xyz", "xyzzzz", "xyabad", "xyyy", "abc", "abbccc", "xyx", "xyxer", "a"]

# 자동 완성 대상 문자열
string = "xy"

status = ["Not found", "Found"]

# Trie 클래스 인스턴스 생성
trie = Trie()

# 문자열 목록으로 트라이 생성
trie.create_trie(strings)

# 입력 문자열에 대한 자동 완성 결과 가져오기
result = trie.show_auto_completion(string)

if result == -1 or result == 0:
    print("No matches")

실행 결과

위 코드를 실행하면 다음과 같이 접두사 "xy"로 시작하는 모든 단어가 출력됩니다.

xyz
xyzzzz
xyabad
xyyy
xyx
xyxer

동작 원리 요약

이 코드의 핵심은 두 단계로 나눌 수 있습니다. 먼저 show_auto_completion 메서드가 사용자가 입력한 접두사까지 트라이를 따라 내려가 해당 위치의 노드를 찾습니다. 이후 matches 메서드가 재귀적으로 그 노드의 모든 하위 분기를 순회하면서 단어가 끝나는 지점(last_node)을 만날 때마다 완성된 단어를 결과 리스트에 추가합니다. 덕분에 공통 접두사를 공유하는 단어들을 빠르게 찾아낼 수 있으며, 이것이 바로 트라이가 검색 자동 완성에 널리 사용되는 이유입니다.