[프로그래머스] 가사 검색 - Trie

처음 접해보는 구조의 문제였고 효율성 테스트까지 상당히 어려웠다..

 

프로그래머스 가사 검색

코딩테스트 연습 - 가사 검색

 

문제는

노래 가사에 사용된 단어들 중, 특정 키워드가 포함된 문자가 몇 개 있는지 찾아내는 문제였다.
여기서 쿼리가 주어지고 그 쿼리에 맞는 단어가 몇개 있는지 찾아내는 것이다.
여기서 쿼리에서는 젤 앞이나 뒤에 "?"라는 문자가 포함될 수 있다. "?"의 개수는 자유롭다.
단, 가운데에 "?"가 포함되는 경우는 없고 앞 뒤에 전부 포함되는 경우도 없다.

그렇다면, 쿼리가 주어지고 각 쿼리당 일치하는 가사가 몇 개나 있는지 찾는 문제였다.

 

가장 처음에 생각했던 방식은 단순하게 O(n^3)으로 풀 수 있는 방식이었다.

그냥 모든 쿼리마다 단어들을 하나씩 비교하고 그 단어마다 앞자리부터 뒷자리부터 단순 비교하는 방식이었다.

당연히 효율성 테스트가 있기 때문에 실패....

 

여기서 처음 접해보는 문제라 많은 생각을 하다... 결국 검색을 통해 자료구조를 알았다.

바로 Trie라는 자료구조였다.

 

이 구조는 관련 검색어를 찾는 알고리즘에 많이 쓰인다고 한다.

기본적으로 Tree로 사용한 구조였다. 말로 설명하기 어려운데, 그림으로 보면 다음과 같다.

보면 각 자리수 별로 Tree로 만들어져서 root로부터 뻗어나오는 것이 나와있다.

 

즉, 각 자리수별로 같은 글자가 있을 시 Tree에 잡아넣고 이를 저장하여 다음에 검색어를 입력할 때, 해당 글자를 기준으로 Tree을 타고 다음 글자를 추천하는 방식입니다.

이렇게 접근을 하게 되면, 어떤 Tree속에서 시간복잡도가 글자수의 길이가 m이라고 한다면 O(m)으로 해결이 가능하다. 

 

이를 이용해 이번 문제를 해결하였다.

그러나 여기서 하나 더 고려를 해야하는 것이 있었다. 

바로 역으로 검색이 가능해야했기 때문에, 역으로 만든 Trie 구조도 필요했다. 그리고 글자 수별로 만들면 쓸모없는 탐색을 수행하지 않을 것 같아서 글자수별로도 만들었다.

 

우선 Tree 구조이기 때문에 각 Node들을 구성할 클래스를 만들었다.

class Node {
    var value: String
    var count: Int
    var child: [String: Nodes]
    
    var isEmpry: Bool {
        return child.isEmpty
    }
    
    init(value: String) {
        self.value = value
        count = 0
        child = [:]
    }
}

 

다음으로 이를 활용해서 Trie 자료구조를 문제에 맞는 메소드를 구성해서 만들었다.

class Trie {
    var root: Node
    
    init() {
        root = Node(value: "")
    }
    
    func insert(_ word: String) {
        var curNode = root
        
        for char in word {
            curNode.count += 1
            if curNode.child[String(char)] == nil {
                curNode.child[String(char)] = Nodes(value: String(char))
            }
            
            curNode = curNode.child[String(char)]!
        }
    }
    
    func search(_ query: String) -> Int {        
        var curNode = root
        for char in query {
            if char == "?" { return curNode.count }
            if curNode.child[String(char)] == nil { return 0 }
            
            curNode = curNode.child[String(char)]!
        }
        
        return curNode.count
    }
}

문제에서 요구사항이 각 노드에 맞는 결과를 return해주는 것이다. 그렇기 때문에 insert을 할 당시부터 해당 노드에 아래에 몇개의 자식이 있는지 저장을 하고 검색을 할 때는 바로 count을 리턴해서 값을 알 수 있게 해주었다.

 

count을 사용하지 않게 되면, 노드를 끝까지 따라가서 개수를 찾아야하기 때문에 비효율적이기 때문이었다.

여기까지만 구현했으면 이제 문제는 어렵지 않았다.

 

글자 수에 따라 Trie을 만들어 주는데 이 때, 역순과 정방향순으로 2가지를 만들면된다.

그리고 Query을 해당 Trie에 넣어주고 count 값을 리턴하면 된다.

 

전체코드는

import Foundation

func solution(_ words:[String], _ queries:[String]) -> [Int] {
    var dic: [Int: Trie] = [:]
    var reverseDic: [Int: Trie] = [:]

    words.forEach { word in
        if dic[word.count] == nil {
            dic[word.count] = Trie()

            reverseDic[word.count] = Trie()
        }

        dic[word.count]?.insert(word)
        reverseDic[word.count]?.insert(String(word.reversed()))
    }

    var answer: [Int] = []
    queries.forEach { query in
        if query[query.startIndex] == "?" {
            // 접두사에 있는 경우 뒤에서부터 탐색해야함 => reverseDic 활용
            if let reverseTrie = reverseDic[query.count] {
                let reverseQuery = String(query.reversed())
                answer.append(reverseTrie.search(reverseQuery))
            } else {
                answer.append(0)
            }
        } else {
            if let trie = dic[query.count] {
                answer.append(trie.search(query))
            } else {
                answer.append(0)
            }
        }
    }

    return answer
}

class Trie {
    var root: Nodes
    
    init() {
        root = Nodes(value: "")
    }
    
    func insert(_ word: String) {
        var curNode = root
        
        for char in word {
            curNode.count += 1
            if curNode.child[String(char)] == nil {
                curNode.child[String(char)] = Nodes(value: String(char))
            }
            
            curNode = curNode.child[String(char)]!
        }
    }
    
    func search(_ query: String) -> Int {        
        var curNode = root
        for char in query {
            if char == "?" { return curNode.count }
            if curNode.child[String(char)] == nil { return 0 }
            
            curNode = curNode.child[String(char)]!
        }
        
        return curNode.count
    }
}

class Nodes {
    var value: String
    var count: Int
    var child: [String: Nodes]
    
    var isEmpry: Bool {
        return child.isEmpty
    }
    
    init(value: String) {
        self.value = value
        count = 0
        child = [:]
    }
}

 

 

정말 어려운 문제였다.... 처음 접해보는 유형이었는데 뭔가 자주 사용할 것 같아서 좋은 문제였던 것 같은 느낌?!