python 리스트를 이용해 trie 구현하면서 개념 익히기

python에서는 trie가 class로 구현된 경우가 많아 꺼렸는데... 배열로 구현하는 법을 익혀서 기록해보기

 

https://deepdata.tistory.com/659

문자열 자료구조 Trie 알고리즘 기본기 배우기

 

https://blog.encrypted.gg/1059

[실전 알고리즘] 0x1F강 - 트라이

 

 

1. trie

 

문자열의 집합은 단순하게 배열 [apple, apply, black, ban, banana, tag, tail, tall]로 저장하거나...

 

{'apple':1, 'apply':1, 'black':1, ...}로 저장하거나...

 

그런데 효율적으로 처리하기 위해 문자열의 집합을 다음과 같이 트리처럼 저장하는 자료구조를 생각할수 있다.

 

 

 

문자열을 하나씩 보면서.. 문자열 내의 문자 하나하나를 하나의 노드에 저장해놓는다

 

루트부터 시작해서 노드를 따라가면서 문자를 하나씩 읽어가다가, 초록색 노드에 도달하면 하나의 문자열을 얻는다.

 

그렇기 때문에 자명하게도 문자열 s 하나를 삽입하는데 O(len(s)), 탐색하는데 O(len(s)), 삭제하는데 O(len(s))이다.

 

 

 

2. 리스트를 이용한 trie의 문자열 삽입

 

class로는 구현해보아서

 

https://deepdata.tistory.com/659

문자열 자료구조 Trie 알고리즘 기본기 배우기

 

 

그런데 class 자체가 알고리즘 문제 풀이에 안맞다고 해야하나? 그런게 있어서 꺼려왔는데

 

배열로 구현하는 방법이 잘 설명되어 있어서 따라해보았다

 

APPLE이라는 단어를 삽입하고 싶은데 트라이의 root를 0번이라고 하면... root의 자식은 1번이다. 

 

먼저 A가 1번에 존재하는지 검사해보고 존재하지 않는다면 A를 1번 노드에 삽입하고

 

다음에 삽입해야할 노드는 2번 노드인데, 다음에 넣을 문자는 P이다.

 

2번 노드에 P가 없기 때문에 P를 2번 노드에 삽입해준다.

 

 

 

이런식으로, 다음에 삽입할 노드는 3번이고 3번에는 P가 없으므로 P를 3번에 삽입하고

 

다음에 삽입할 노드는 4번이고 4번에는 L이 없으므로 L을 4번에 삽입하고

 

다음에 삽입할 노드는 5번이고 5번에는 E가 없으므로 E를 5번에 삽입한다

 

 

 

 

E를 끝으로 APPLE을 전부 삽입했으므로, 5번 노드는 끝났다는 표시를 해준다.

 

다음 APPLY라는 단어를 삽입하고 싶은데, 다시 root인 0번 노드부터 시작해서 

 

0번 노드의 자식 중에 A가 존재하므로 A를 삽입할 노드를 만들지 않고 A의 노드 번호인 1번으로 이동

 

다음 P를 삽입해야하는데 1번 노드의 자식인 2번 노드에 P가 있으므로 2번 노드로 이동

 

...

 

이런 식으로 A,P,P,L까지 4번 노드까지 이동하게 된다.

 

그런데 다음 Y를 삽입할려고 하는데 L의 자식 중에 Y가 존재하지 않으므로 지금까지 사용한 노드 번호 5번이 아닌,

 

새로운 노드 번호 6번에 Y를 삽입해준다.

 

 

 

APPLY를 모두 삽입했으니, 6번 노드에는 끝났다는 표시를 해준다.

 

다음 단어로 BLACK을 삽입하고 싶다.

 

다시 ROOT인 0번 노드부터 시작하는데, 0번 노드의 자식 중에 B가 존재하지 않으므로,

 

0번 노드의 자식으로 지금까지 사용하지 않은 번호인 7번 노드에 B를 삽입하게 된다.

 

 

 

그리고 다음 L을 삽입할때 7번 노드의 자식에 8번 노드로 L을 삽입하고

 

A를 삽입할때, 8번 노드의 자식에 9번 노드로 A를 삽입하고

 

C를 삽입할 때, 9번 노드의 자식에 10번 노드로 C를 삽입하고

 

K를 삽입할때, 10번 노드의 자식에 11번 노드로 K를 삽입하게 된다.

 

그리고 BLACK를 모두 삽입했으니 11번 노드에 끝났다는 표시를 해준다.

 

 

 

 

만약 이걸 파이썬의 리스트로 저장한다고 하면... 각 노드 번호는 리스트의 인덱스를 나타내고, 해당 인덱스에 문자를 저장한다.

 

trie = [-1, A, P, P, L, E, L, B, L, A, C, K]

 

end = [0,0,0,0,0,1,1,0,0,0,0,1]과 같이 끝났다는 표시를 해주는 배열도 필요하다.

 

 

3. 구현 예시

 

trie를 2차원 배열로 선언하는데,

 

trie[가능한 모든 문자의 개수][문자의 종류]로 선언한다.

 

문자의 종류는 알파벳 소문자만 쓴다면 26이겠지

 

가능한 모든 문자의 개수는 APPLE, APPLY, BLACK라면 각 문자들이 노드가 될수 있으므로 문자열 길이의 합으로 선언했다

 

그런 다음 끝을 표현해주는 end 배열도 선언해줘야한다.

 

L = 0 #가능한 모든 문자의 개수, 가능한 노드 번호의 최댓값

for s in S:
    
    L += len(s)

#trie[문자열 집합 내에 존재하는 모든 문자의 최대 개수][문자의 종류]
trie = [[-1]*26 for _ in range(L+1)]
end = [0]*(L+1)

 

 

문자열 집합 내의 문자열 s를 순회하는데, root번호 0번부터 current_node로 시작해서...

 

current_node의 자식은 trie[current_node]이다.

 

현재 삽입해야할 문자는 s를 순회하면서 문자 c를 ord로 바꿔서 숫자에 대응시킨 다음

 

trie[current_node][ord(c)]가 -1인지 아닌지 검사해본다.

 

trie[current_node][ord(c)]가 -1이 아니라면 current_node의 자식으로 c가 존재한다는 것이다.

 

-1이라면 자식에 c가 존재하지 않으므로 아직 사용하지 않은 노드 번호인 num을 대응시켜주고, num += 1을 해준다.

 

그리고 나서 다음 노드 번호인 trie[current_node][ord(c)]로 이동해준다.

 

모든 문자를 삽입하면 end[current_node] = 1로 끝났다는 것을 체크한다.

 

#배열을 이용한 trie의 단어 삽입

n = int(input())

S = input().split()

L = 0 #가능한 모든 문자의 개수, 가능한 노드 번호의 최댓값

for s in S:
    
    L += len(s)

#trie[문자열 집합 내에 존재하는 모든 문자의 최대 개수][문자의 종류]
trie = [[-1]*26 for _ in range(L+1)]
end = [0]*(L+1)

num = 1 #아직까지 사용하지 않은 노드 번호(삽입할때 사용할 노드 번호)

for s in S:
    
    current_node = 0 #현재 노드 번호

    for c in s:
        
        v = ord(c) - ord('a') #현재 문자
        
        #현재 노드 번호 current_node의 자식 trie[current_node]
        #trie[current_node][v] = -1이라는 것은,
        #current_node의 자식으로 현재 문자 v가 없다는 것
        if trie[current_node][v] == -1:
            
            #현재 노드의 자식에 현재 문자 v가 없으면
            #아직 사용하지 않은 노드 번호 num을 만들어주고, num += 1해서 다음에 사용
            trie[current_node][v] = num
            num += 1
        
        #v를 삽입하고, v의 노드 번호로 이동
        current_node = trie[current_node][v]
    
    end[current_node] = 1 #모든 문자를 삽입했으면, 끝났다는 표시
            
print(trie)
print(end)

 

 

 

 

구현을 보면 알겠지만 기본적으로 메모리를 상당히 많이 사용하게 된다.

 

삽입하는데 가능한 모든 문자의 개수가 문자열 집합 내 모든 문자열의 길이 합 * 문자의 종류만큼 공간을 차지하는데

 

알파벳 소문자만 쓴다면 26개라서.. 문자열 길이 합이 100만정도라면 2600만이라 조금 힘들수 있을것 같다

 

그러다보니 class나 dict를 사용해서 메모리를 절약하는듯??

 

 

4. 문자열 탐색

 

삽입을 구현하는 방법을 이해했다면, 어떤 문자열이 존재하는지 탐색하는 것은 매우 간단하다.

 

root 노드 0번 = current_node부터 시작해서,

 

문자열 s를 순회하여 문자 c에 대해, current_node의 자식 trie[current_node]에 ord(c)가 존재하는지 체크

 

존재하지 않는다면 문자열 s는 존재하지 않으니까 반복문을 바로 탈출

 

존재한다면 다음 자식 노드 current_node = trie[current_node][ord(c)]로 이동

 

반복문이 끝나고, 마지막으로 도달한 노드 current_node가 실제로 끝난 노드인지 체크해줘야한다.

 

end[current_node] = 1이라면, 문자열 s가 존재하는 것이고 0이라면 존재하지 않는 것이다.

 

s = 'black'

current_node = 0 #root인 0번 노드부터 시작

no = False

for c in s:
    
    v = ord(c) - ord('a')
    
    #현재 노드의 자식 trie[current_node]중 문자 c가 존재하는지 체크
    if trie[current_node][v] == -1:
        no = True
        break #존재하지 않는다면 바로 반복문 탈출
    
    #존재한다면 다음 자식 노드 trie[current_node][v]로 이동
    current_node = trie[current_node][v]

if no:
    
    print('NO')

else:
    
    #마지막으로 도달한 노드 current_node가 실제로 끝난 노드인지 검사
    if end[current_node] == 1:
    
        print('YES')
    
    else:
        
        print('NO')

YES

 

 

주의할 점은 마지막으로 도달한 노드가 실제로 끝난 노드인지 검사해줘야한다.

 

APP가 존재하는지 검사할때, root인 0번부터 시작해서...

 

0번 노드의 자식 중 A가 1번 노드에 존재하므로 1번 노드로 이동

 

1번 노드의 자식 중 P가 2번 노드에 존재하므로 2번 노드로 이동

 

2번 노드의 자식 중 P가 3번 노드에 존재하므로 3번 노드로 이동

 

이러면 반복문이 끝났는데, 실제로 trie에는 APPLE, APPLY, BLACK밖에 없다.

 

그러니까 3번 노드는 끝난 노드가 아니기 때문에 APP가 없다고 답해야한다.

 

 

 

5. 문자열 제거

 

문자열을 제거할때도 삽입과 마찬가지로 root부터 시작한다음,

 

문자열을 순회하면서 해당 문자가 현재 노드의 자식에 존재한다면 자식 노드로 이동해준다.

 

물론 자식에 존재하지 않는다면 해당 문자열 자체가 존재하지 않으므로 삭제할 필요도 없다.

 

자식 노드로 이동하다가 노드의 끝에 도달했다면, 해당 노드가 "끝난 노드"에서 "끝나지 않은 노드"로 바꿔주기만 하면 된다

 

 

s = 'black'

current_node = 0

find = True

for c in s:
    
    v = ord(c) - ord('a')

    if trie[current_node][v] == -1:
        
        find = False
        break
    
    current_node = trie[current_node][v]

if find:
    
    end[current_node] = 0

 

 

문자열을 삭제한다고 해서 다음과 같이 노드 자체를 삭제하면 조금 곤란하다

 

APPLE을 삭제할려고 할때, 1번노드 A를 삭제해버리고, 2번 노드 P를 삭제해버리고

 

3번 노드 P를 삭제해버리고, 4번 노드 L을 삭제해버리고 5번 노드 E를 삭제해버리고..

 

이러면 APPLY도 없어져버려서... 트라이 구조 자체를 망치게 된다.

 

그러다보니 트라이는 문자열 삭제가 자주 발생하는 경우에는 효율적이지 않다.

 

또한 정점이 남아있다보니 애초에 삭제 자체가 효율적이라고 보기 어렵다