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

문제 풀이

문제 설명은 해당 링크를 참조하자.

각 쿼리(n개), 각 단어(m개)를 매칭하여 비교하려고 하면 시간 복잡도가 O(nm)이어서 호율성 테스트를 통과할 수가 없다. 따라서 단어를 최대한 그룹화할 필요가 있는데, 이때 사용할 수 있는 것이 바로 Trie라는 자료구조이다.

'f'와 'k'보다 위에 있는 root로는 아무 문자도 들어있지 않은 빈 node를 사용한다.

위 그림과 같이 앞에서부터 글자들을 비교하여 같은 문자들의 경우 공통된 node로 해결하고, 달라지기 시작하는 문자에서부터 다른 node로 처리하는 것이다. 그리고 각 node마다 문자열에 길이를 key 값으로 입력하면 해당 길이를 가지는 문자열의 수를 반환하는 map을 설정한다.

예를 들면 왼쪽의 'o' node에 있는 map에 key값으로 5를 입력하면 맨 앞에 'fro'라는 부분 문자열로 시작하면서 길이가 5인 문자열의 개수를 알 수가 있는 것이다.

이러한 Trie를 이용하면 물음표가 뒤에 달려있는 Query를 해결할 수 있다.

물음표가 앞에 달려 잇는 Query 또한 Query를 reverse한 다음 word들의 역문자열들을 저장한 reversedTrie를 이용하여 해당되는 문자열의 개수를 구할 수 있다.

import java.util.HashMap;

class Solution {
    class Trie{
        char root; // node에 해당
        HashMap<Character, Trie> subTries = new HashMap<>();
        HashMap<Integer, Integer> lenMap = new HashMap<>();
        // root에서부터 현재 node까지의 단어들을 순서대로 가지는 문자열들 중
        // key 값에 해당하는 길이를 가지는 문자열의 개수

        public Trie(char c){
            root = c;
        }

        public void insertWord(String word){
            Trie current = this;

            for(char c : word.toCharArray()){
                // 처음으로 기존 string로부터 없는 문자가 나오는 경우
                if(!current.subTries.containsKey(c))
                    current.subTries.put(c, new Trie(c));

                if(!current.lenMap.containsKey(word.length()))
                    current.lenMap.put(word.length(), 0);
                current.lenMap.put(word.length(), 
                                   current.lenMap.get(word.length()) + 1);

                current = current.subTries.get(c);
            }
        }

        public int count(String query){
            Trie current = this;

            for(char c: query.toCharArray()){
                if(c != '?'){
                    // 해당 문자를 해당 자리에 가진 문자열이 없는 경우
                    // 0 반환
                    if(!current.subTries.containsKey(c)) return 0;
                    current = current.subTries.get(c);
                }else{
                    // 처음으로 물음표가 나온 경우에는 문자열 길이만 체크하면 됨
                    return current.lenMap.getOrDefault(query.length(), 0);
                }
            }

            return 0;
        }
    }

    public String reverse(String str){
        return new StringBuilder(str).reverse().toString();
    }

    public int[] solution(String[] words, String[] queries) {
        Trie trie = new Trie('\0');
        // 각 단어를 저장하는 trie 설정

        for(String word: words)
            trie.insertWord(word);


        Trie reversedTrie = new Trie('\0');
        // 역순으로 나열한 단어들을 저장하는 trie 설정

        for(String word: words)
            reversedTrie.insertWord(reverse(word));

        int[] answer = new int[queries.length];
        for(int i = 0;i < queries.length;i++){
            String query = queries[i];
            if(query.charAt(0) != '?')
                answer[i] = trie.count(query);
                // 물음표가 뒤에 있는 query를 처리
            else answer[i] = reversedTrie.count(reverse(query));
            // 물음표가 앞에 있는 query를 처리
        }

        return answer;
    }
}