[알고리즘] Trie(트라이) 알고리즘 및 게임 닉네임(16934) 파이썬 풀이

2023. 8. 13. 14:19알고리즘

※ 사용언어: 파이썬

※ 트라이 알고리즘 설명 아래 16934 게임닉네임 백준 파이썬 풀이 코드가 있습니다.

 

 

개념

문자열을 효율적으로 탐색하기 위해, 문자열을 트리형태로 저장한 자료구조이다. 다른 이름으로는 radix tree 혹은 prefix tree가 있다. Trie 알고리즘은 크게 문자열을 트리형태로 저장하는 삽입과 문자열을 검색이라는 두 부분으로 나눌 수 있다.

 

 

Trie 알고리즘 - 삽입

"apple, car, chat"이라는 문자열이 있다고 가정하겠다. 해당 문자열을 트리 구조로 표현하면 다음과 같다. (여기서 "head"는 트리의 root 노드라고 생각하면 된다.)

트리 구조로 표현한 문자열

해당 트리 구조를 기억하기 위해서는 각 노드들은 3가지의 변숫값을 기억하고 있어야 한다.

  • key: 현재 노드
  • data: 최종 문자열 (leaf node가 가지는 값)
  • children: 노드가 가지고 있는 자식

예시를 통해 구체적인 설명을 하기 위해 head아래 있는 C 부분만 따로 가져오겠다. 처음 car를 저장하는 과정은 다음과 같다. c와 a는 각각 아래 자식 노드가 있기 때문에 children 값을 가지지만, r 같은 경우 자신이 leaf node이기 때문에 위에 있는 모든 노드들의 문자를 합친 최종 문자열 값을 data에 저장한다.

 

car를 저장한 뒤, chat을 저장하게 되면 다음과 같이 테이블이 바뀌게 된다.

c는 이제 a와 h를 자식으로 가지게 되고, 새롭게 들어온 h와 a는 각각 아래 있는 노드를 children에 저장한다. 마지막 t는 leaf node로 최종 문자열을 기억하게 된다.

 

위의 두 과정을 한 그림으로 보면 다음과 같다.

코드로 해당 과정을 구현하게 되면 다음과 같다.

우선 3가지의 변숫값을 기억할 class를 하나 만든다.

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

변수값을 기억할 class를 이용해 각 노드에 변숫값들을 저장시켜 준다.

class Trie(object):
    def __init__(self):
        self.head = Node(None)

    def insert(self, string):
        # 현재 노드 확인
        now = self.head
        for char in string:
            # 자식 노드 중에 중복 문자가 없는 경우
            # 자식 노드에 문자 추가
            if char not in now.children:
                now.children[char] = Node(char)
            # 같은 문자가 있으면 문자를 추가하지 않고
            # 자식 문자로 내려가기
            now = now.children[char]
        # 마지막 도착 == leaf node
        # data에 최종 문자 추가
        now.data = string

 

 

Trie 알고리즘 - 검색

검색을 할 문자열을 받으면, 문자열에 있는 문자를 순차적으로 순회하며, 해당 문자가 children을 가지고 있는지 검사한다. children 노드가 없을 때까지 아래로 파고들며, 마지막 leaf node에 도착하게 되면 leaf node가 data를 가지고 있는지 확인한다. leaf node가 data를 가지고 있으면 해당 문자열이 트리 안에 저장된 것으로 판단할 수 있다.

(코드상에서는 data가 있는 경우 data를 출력하게 했고, 아닌 경우에는 False를 출력하게 했다.)

def search(self, string):
    now = self.head
    for char in string:
        if char in now.children:
            now = now.children[char]
        else:
        	# 해당 문자열이 없는 경우
            return False
    # 해당 문자열이 있는 경우
    if now.data != None:
        return now.data
# main 영역
test = Trie()
test.insert("apple")
test.insert("car")
test.insert("chat")

print(test.search("car"))
print(test.search("abc"))

 

 

Trie 알고리즘 

삽입 알고리즘과 검색 알고리즘을 모두 합친 전체 알고리즘은 다음과 같다.

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

class Trie(object):
    def __init__(self):
        self.head = Node(None)

    def insert(self, string):
        # 현재 노드 확인
        now = self.head
        for char in string:
            # 자식 노드 중에 중복 문자가 없는 경우
            # 자식 노드에 문자 추가
            if char not in now.children:
                now.children[char] = Node(char)
            # 같은 문자가 있으면 문자를 추가하지 않고
            # 자식 문자로 내려가기
            now = now.children[char]
        # 마지막 도착 == leaf node
        # data에 최종 문자 추가
        now.data = string

    def search(self, string):
        now = self.head
        for char in string:
            if char in now.children:
                now = now.children[char]
            else:
                # 해당 문자열이 없는 경우
                return False
        # 해당 문자열이 있는 경우
        if now.data != None:
            return now.data


test = Trie()
test.insert("apple")
test.insert("car")
test.insert("chat")

print(test.search("car"))
print(test.search("abc"))

# result
# car
# False

 

 

백준 - 게임닉네임(16934)

트라이 알고리즘을 사용하는 문제이다. 문제에서 트라이 알고리즘을 사용하라고 접두사를 사용하여 닉네임을 저장한다고 업급하고있다.

위에 설명한 트라이 알고리즘에서 추가적으로 추가해야 하는 부분들이 있다.

  • 가장 길이가 짧은 유저 별칭 출력
  • 같은 닉네임으로 여려 명이 가입하여 더는 별칭을 부여할 수 없는 경우, 닉네임 뒤에 번호를 붙여 저장.
    • 예를 들어 backjoon이라는 이름에 더는 별칭을 붙일 수 없는데, backjoon이라는 이름을 누군가 게임 닉네임으로 사용한 경우 별칭은 backjoon2가 된다. 그 이후 다시 backjoon이란 이름을 사용한 누군가가 나타나면 그 사람은 backjoon3이라는 별칭을 가지게 된다.
import sys
from collections import defaultdict

N = int(sys.stdin.readline())
name = []
for _ in range(N):
    name.append(sys.stdin.readline().rstrip())


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

class Trie(object):
    # 루트 설정
    def __init__(self):
        self.head = Node(None)

    def insert(self, string):
        now = self.head
        for char in string:
            if char not in now.children:
                now.children[char] = Node(char)
            now = now.children[char]
        now.date = string
        # 중복 닉네임 뒤에 번호 붙여야하니 번호 지정
        same_nick[string] += 1

    def search(self, string):
        now = self.head
        # 가장 짧은 별칭을 저장할 result 배열 선언
        result = ''
        for char in string:
            result += char
            if char in now.children:
                now = now.children[char]
            else:
                # 검색을 먼저 진행하여, 처음  baekjoon를 검색할 때
                # 저장된 닉네임이 없는 상태이기에 b는 자식 노드가 없습니다.
                # else문으로 넘어오게 되고, 바로 return
                return result
        # 해당 닉네임이 저장되어 있어
        # 짧은 별칭을 지정할 수 없는 경우입니다.
        if now.date != None:
            # same_nick 딕셔너리 초기값이 0부터 시작해서 1 더해줌
            result += str(same_nick[string] + 1)
        return result


check = Trie()
# 중복 이름이 몇개 있는지 저장할 딕셔너리
# 초기화를 위해 defaultdict로 선언
same_nick = defaultdict(int)
for i in range(N):
    print(check.search(name[i]))
    check.insert(name[i])