Trie(트라이) 자료구조 파이썬

 

 

Node 선언

class Node(object) :
    def __init__(self,data):
        self.data = data
        self.count = 0
        self.child = {}

Node 선언 클래스를 만들어준다. Node는 각각 data, count, child를 가지고 있다.

data : 현재 노드가 어떤 문자를 가지는지를 보여준다. Node('a')라고 선언할 경우, 자동으로 초기화하는데 이 때 data에는 'a'가 들어가게 된다.

count : 현재 노드를 지나간 문자열이 몇개인지 알려주는 변수. 해당 노드를 지나가는 문자열이 하나씩 생길 때 마다, count를 증가시킨다.

child : 현재 노드의 자식노드를 나타낸다. 딕셔너리 형태로 선언하고, 딕셔너리의 key는 data명이, value에는 다음 Node가 들어온다.

 

Trie 구조

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

    def insert(self, string):
        cur = self.head
        cur.count +=1

        for c in string :
            if c not in cur.child :
                cur.child[c] = Node(c)
            cur = cur.child[c]
            cur.count += 1


    def count(self,prefix):
        cur = self.head
        
        for c in prefix :
            if c not in cur.child :
                return 0
            cur = cur.child[c]

        return cur.count

처음 Trie를 만들 경우, 초기화는 Head Node를 만드는 것으로 시작한다. 당연하지만, Head Node는 아무 값도 가지지 않도록 'None'으로 생성한다.

Insert 함수는 Trie에 문자열을 집어넣는 함수다. 다음과 같은 로직으로 진행된다.
1. 객체의 Head를 현재 Node로 설정한다.
2. 각 문자열에 대해 다음과 같은 과정을 반복한다. 
 2-1. 현재 문자열이 현재 노드의 자식 노드가 아니라면, 자식 노드를 추가한다.
 2-2. 현재 노드를 자식 노드로 업데이트 한다. 
 2-3. 방문한 노드의 count를 올려준다.

Insert 과정을 통해 'aaa', 'abc', 'abcd', 'bcd'를 넣는 것을 도식화하면 다음과 같다.

1. 'aaa'를 넣는다. Head Node에 자식노드로 a가 추가되며, 각 노드마다 count가 1씩 증가한다.

2. 'abc'를 넣는다. Head와 Head에 연결된 자식 노드 'a'를 지나간 문자열은 2번이기 때문에 Count가 2로 되었다.

3. 'abcd'를 업데이트 한다. a-b-c는 2번 지나왔으니 각 노드의 Count가 2가 되고, 마지막 노드 'C'에 자식노드 'D'가 다시 추가된다.

4. 'bcd'를 추가한다. Head Node에는 자식노드로 'b'가 없기 때문에 새로운 자식노드가 만들어져야 한다.

 

참고할 부분

무지성 코더인 나는 처음에 딕셔너리가 같은 이름인 'a'를 가지면, 서로 연결되는거 아닌가? 라는 생각을 가지고 위의 Node Class의 활용을 잘 이해하지 못했다. 이 때는 Class의 속성에 대해서 전혀 모르고 있었기 때문이다.

쉽게 생각하면 리스트를 생각하면 된다. 예를 들어 a = [{'a' : '1'}], b = [{'a' : '1' }] 이렇게 선언을 한다면 당연히 a,b 안에 있는 것은 같은 값을 가리키지만 실제로 저장된 주소는 다르다.

다음과 같이 이해를 해보면 어떨까?

A,B 라는 Trie 구조를 각각 만든다. 그럼 각 Head Node의 이름은 A.head, B.head가 될 것이다. 그리고 각각의 Head Node는 Child라는 Dictionary를 가진다. 즉, A.Head.Child = {}, B.Head.Child = {}가 된다. 이 딕셔너리에 'a'라는 Key를 가진 Node를 각각 추가한다. 이들은 결코 다르다. 왜냐하면 각기 다른 딕셔너리에서 각각 선언되었기 때문이다. 

Count 함수는 Trie에서 해당 문자열을 만족하는 문자열이 몇개 있는지를 확인하는 함수다. 다음과 같은 로직으로 진행된다.

1. 객체의 Head를 현재 Node로 설정한다.
2. 현재 문자열이 자식 노드에 있는지 확인한다. 
   2-1. 자식 노드에 없다면, 이런 문자열이 현 Trie 구조에 존재하지 않기 때문에 0을 돌려준다.
   2-2. 자식 노드에 있다면, 현재 노드를 자식 노드로 업데이트 해주고 반복문을 반복한다.
3. 반복문을 끝까지 돌았다면, 현재 노드에 있는 Count 값을 리턴해준다.