[프로그래머스 > Lv4] 가사 검색

프로그래머스

 

1. 문제

[입력]

• 가사에 사용된 모든 단어들이 담긴 배열 words
• 찾고자 하는 키워드가 담긴 배열 queries

 

[결과]

각 키워드별로 매치된 단어가 몇 개인지 순서대로 배열에 담아 return

 

2. 풀이

처음은 완전 탐색 방식으로 간단하게 구현하였으나, 효율성 2번에서 계속 시간초과가 발생하였습니다. 😭

결국 알고리즘으로 풀어보기로 결정해서, 사용한 알고리즘은 바로 Trie입니다. 🤗

 

2.1 Trie Class

Trie에는 2가지 메서드를 정의해였습니다.

• 단어 추가
• query와 일치하는 단어 개수 탐색

class Trie {
    var count = 0
    private val children = hashMapOf<Char, Trie>()

    fun add(word: String, index: Int, direction: Int) {
        if (index<0 || word.length==index) {
            count = 1
            return
        }
        val cur = word[index]
        if (cur !in children) children[cur] = Trie()
        children[cur]!!.add(word, index + direction, direction)
    }

    fun count(query: String, index: Int): Int {
        if (index == query.length) return count

        if (children.size == 0) return 0
        if (query[index] != '?') {
            if (query[index] !in children) return 0
            return children[query[index]]!!.count(query, index+1)
        }
        return children.values.sumOf { it.count(query,index+1) }
    }
}

 

이를 이용해 입력된 word들의 trie tree를 만들어주면 됩니다. 이곳에서 중요한 것은!

🌹 word의 앞에서부터 시작하는 trie tree와 word의 뒤에서부터 시작하는 trie tree

위 2가지가 필요하다는 점입니다. 👍🏻

 

👇🏻 요렇게!

val trie = Trie()
val reverseTrie = Trie()
words.forEach {
    trie.add(it, 0, 1)
    reverseTrie.add(it, it.length - 1, -1)
}

 

이제 이 tree들을 이용해, 각 query에 해당하는 단어의 개수를 구해주면 됩니다.

여기서 추가적으로 검색 키워드 제한사항에 있는

검색 키워드는 중복될 수 있습니다.

라는 점을 이용해 실행 시간을 줄이기 위해 HashMap을 사용해주도록 하겠습니다.🔥

 

3. 코드

Kotlin

class Solution {
    class Trie {
        var count = 0
        private val children = hashMapOf<Char, Trie>()

        fun add(word: String, index: Int, direction: Int) {
            if (index<0 || word.length==index) {
                count = 1
                return
            }
            val cur = word[index]
            if (cur !in children) children[cur] = Trie()
            children[cur]!!.add(word, index + direction, direction)
        }

        fun count(query: String, index: Int): Int {
            if (index == query.length) return count
            
            if (children.size == 0) return 0
            if (query[index] != '?') {
                if (query[index] !in children) return 0
                return children[query[index]]!!.count(query, index+1)
            }
            return children.values.sumOf { it.count(query,index+1) }
        }
    }

    fun solution(words: Array<String>, queries: Array<String>): IntArray {
        val n = queries.size
        val answer = IntArray(n) { 0 }

        val trie = Trie()
        val reverseTrie = Trie()
        words.forEach {
            trie.add(it, 0, 1)
            reverseTrie.add(it, it.length - 1, -1)
        }

        val dictionary = hashMapOf<String, Int>()
        for ((index, query) in queries.withIndex()) {
            answer[index] = dictionary[query] ?: run {
                val type = query.startsWith("?").not()
                val count = if (type) trie.count(query, 0) else reverseTrie.count(query.reversed(), 0)
                dictionary[query] = count
                count
            }
        }
        return answer
    }
}