[알고리즘/Swift] BOJ 14725 개미굴 :: 문자열, 트라이 트리 (Trie Tree)

알고리즘 스터디 지난주 주제는 "문자열" 이었습니다. 트라이 트리를 활용한 문제가 인상적이었고, 많이 활용되는 알고리즘이라고 합니다. 이번주 Swift 알고리즘 책에서도 트리 중, 트라이 트리를 공부해서 정리해보려고 합니다!

14725번: 개미굴

 

👀 풀이 방법

이 문제에서는 아래 그림으로 힌트를 줬어요. 처음 문제를 봤을 땐 트라이 트리를 모르는 상태였지만, 아.. 뭔가 트리구조를 활용해야하는 구나를 알게 되었죠! 그래서 찾아보니, 트라이 트리를 활용해서 푸는 문제라는 걸 알게 되었습니다.

 

💡 트라이 트리 (Trie Tree) ?
트라이 트리는 문자열 검색에 많이 사용되는 알고리즘이다. 모서리(Edge) 에 Key 를 지니고, 노드의 모든 자손은 노드에 연관된 문자열의 공통 접두사(Prefix)를 공유한다. Prefix Tree 라고도 부른다.

개미굴 예시 이미지

✔︎ 입·출력

개미굴 문제는 입력된 전체 문자열을 띄어쓰기를 기준으로 노드를 구분해 삽입하고, 이를 시각화해서 출력하는 문제입니다. 

입 / 출력

 

✔︎ Swift 코드

트라이 트리를 적용하기 위해 노드와 트리 클래스를 구현하고, 삽입과 전체 트리를 탐색하는 메소드를 구현하면 됩니다.

import Foundation

class TrieNode {
    var key: String?
    weak var parent: TrieNode?
    var children: [String: TrieNode] = [:]
    var isLeaf: Bool {
        return children.isEmpty
    }
    
    init(key: String? = nil, parent: TrieNode? = nil) {
        self.key = key
        self.parent = parent
    }
}

class TrieTree {
    // 트라이 트리의 루트의 key 는 nil이다.
    let root: TrieNode
    
    init() {
        root = TrieNode()
    }
    
    // 입력 문자열을 통해 노드를 삽입하는 메소드
    func insert(words: [String]) {
        var currentNode = root
        let n = Int(words[0])!
        for word in words[1...n] {
            
            // 단어가 현재 노드의 자식 노드에 이미 존재하는 경우 -> 현재 노드 변경
            if let childNode = currentNode.children[word] {
                currentNode = childNode
            }
            
            // 단어가 현재 노드의 자식 노드에 존재하지 않는 경우 
            // -> 노드 생성, 자식 노드 추가, 현재 노드 변경
            else {
                let newNode = TrieNode(key: word, parent: currentNode)
                currentNode.children[word] = newNode
                currentNode = newNode
            }
        }
    }
    
    // 트라이 트리를 전체적으로 탐색하는 메소드
    func traversal(node: TrieNode, depth: Int) {
        
        // 트리의 계층을 나타내기 위한 출력문 - root를 제외한 모든 노드가 출력된다.
        if let key = node.key {
            var dash: String = ""
            for _ in 0..<depth {
                dash += "--"
            }
            print(dash + key)
        }
        
        // 잎사귀 노드인 경우 재귀적으로 호출된 메소드를 종료한다.
        guard !node.isLeaf else {
            return
        }
        // 노드의 자식들을 사전순으로 탐색하기 위해 딕셔너리를 정렬한다.
        let sortedChildren = node.children.sorted(by: {$0.0 < $1.0})
        
        // 노드의 자식들을 재귀적으로 탐색한다.
        for child in sortedChildren {
            traversal(node: child.value, depth: depth + 1)
        }
    }
}

// 입력과 함수 호출 등
let N = Int(readLine()!)!
let trie = TrieTree()
for _ in 0..<N {
    let line = readLine()!.split(separator: " ").map{String($0)}
    trie.insert(words: line)
}
trie.traversal(node: trie.root, depth: -1)

 

Python3 코드

처음에 Python 으로 풀이를 시작했어서 참고용으로 올려둡니다 :)

class Node:
    # key - 노드의 키, data : 전체 문자열, children : 자식 노드 딕셔너리
    def __init__(self, key, data=None):
        self.key = key
        self.data = data
        self.depth = 0
        self.children = {}

class Trie:
    def __init__(self):
        self.head = Node(None)
    
    def insert(self, string):
        current_node = self.head
        depth = 0
        for s in string.split(' '):
            # 새로운 문자가 노드의 하위에 없는 경우, 딕셔너리를 생성
            if s not in current_node.children:
                current_node.children[s] = Node(s)
            # 현재 노드를 변경
            current_node = current_node.children[s]
            current_node.depth = depth
            #print(current_node.key, current_node.data == None)
            depth += 1
        # 하단 노드에 전체 문자열을 저장
        current_node.data = string
        #print(current_node, current_node.data == None)

    def traversal(self, current_node):
        #current_node = self.head
        if current_node.data:
            return

        children = sorted(current_node.children.keys())
        for c in children:
            depth = current_node.children[c].depth
            print('--'*depth, end='')
            print(c)
            self.traversal(current_node.children[c])

N = int(input())
trie = Trie()
for _ in range(N):
    data = input().split()
    data = ' '.join(data[1:])
    trie.insert(data)

trie.traversal(trie.head)