본문 바로가기

알고리즘

프로그래머스 가사검색 코틀린

반응형
class Solution {
    fun solution(words: Array<String>, queries: Array<String>): IntArray {
        val answer = IntArray(queries.size)
        val root = Trie()
        val tail = Trie() //역으로 탐색하기 위해

        for (word in words) {
            root.insert(word)
            tail.insert(word.reversed()) //역으로 입력
        }
        for (i in queries.indices) {
            if (queries[i][0]!='?') {
                answer[i] = root.find(queries[i])
            } else {
                answer[i] = tail.find(queries[i].reversed())
            }
        }
        return answer
    }

    private class Trie {
        val children: HashMap<Char, Trie?> = HashMap() //알파뱃 개수에 맞추어 배열로해도 됨
        val len: HashMap<Int, Int> = HashMap() //알파뱃 개수에 맞추어 배열로해도 됨(각 trie마다 최대 길이를 저장 ex. frod, frro와 frr?이 있다면 f,r은 2 frr은1이다.)

        fun insert(text: String) {
            var trie = this

            for (ch in text) {
                if (trie.children[ch]==null) {
                    trie.children[ch] = Trie()
                }
                trie.len[text.length ] = trie.len.getOrDefault(text.length, 0) + 1
                trie = trie.children[ch]!!
            }
        }

        fun find(text: String): Int {
            var trie = this
            for (ch in text) {
                if (ch == '?') {
                    return trie.len.getOrDefault(text.length, 0)
                } else if(trie.children[ch]==null) { //다음꺼 참조할 경우가 없을 경우 0을 리턴
                    return 0
                }
                else {
                    trie = trie.children[ch]!!
                }
            }
            return trie.len.getOrDefault(text.length, 0)
        }
    }
}
반응형