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

기본 Trie 문제에서 약간 변형이 된 문제이다. 아이디어를 이해하면 풀 수 있을 것이다. 

접근 방법

HashMap<Integer, Integer> lengthMap 을 운영하여 각 단계 마다 뒤에 length가 몇개 있는지 확인 진행한다.

lengthMap을 운영하면 fro??, fro???인 경우를 구분할 수 있다. 

그렇기 때문에 Trie에 insert 시 아래와 같이 작성해주어 length의 갯수가 몇개 있는지 확인하고 1씩 더해준다. 

int len = str.length();
cur.lengthMap.put(len, cur.lengthMap.getOrDefault(len, 0) +1);

 

search로 찾을 때에는 현재 찾고 있는 str의 length로 HashMap을 접근해준다.

만약, "??" 처럼 queries[]에 동일한 글자수가 없을 경우가 있기 때문에 없으면 0을 호출하기 위해서 아래와 같이 사용해준다.

cur.lengthMap.getOrDefault(str.length(), 0);

 

전체 코드 

import java.util.ArrayList;
import java.util.HashMap;

class Node{
    HashMap<Character, Node> child;
    HashMap<Integer, Integer> lengthMap;
    int cnt;
    public Node(){
        this.child = new HashMap<Character, Node>();
        this.lengthMap = new HashMap<Integer, Integer>();
        this.cnt = 0;
    }
}
class Trie{
    Node root;
    public Trie(){
        this.root = new Node();
    }

    public void insert(String str){
        Node cur = this.root;

        int len = str.length();

        for(int i=0; i<str.length(); i++){
            char c = str.charAt(i);
            // 현재 해당하는 문자가 cur 기준으로 없으면 넣어주고 있으면 넘어감
            cur.cnt++;
            cur.child.putIfAbsent(c, new Node());
            cur.lengthMap.put(len, cur.lengthMap.getOrDefault(len, 0)+1);
            cur = cur.child.get(c); // 해당 위치로 이동함
        }
    }

    public int search(String str){
        Node cur = this.root;

        for(int i=0; i<str.length(); i++){
            char c = str.charAt(i);
            if(c == '?') {
                return cur.lengthMap.getOrDefault(str.length(), 0);
            }else if(cur.child.containsKey(c)){
                cur = cur.child.get(c); // 해당 노드로 이동
            }else{
                return 0;
            }
        }
        return cur.cnt;
    }
}

public class PGM {
    public static int[] solution(String[] words, String[] queries) {
        int[] answer = {};
        Trie trie = new Trie();
        Trie reverseTrie = new Trie();
        for(int i=0; i<words.length; i++){
            String s = words[i];
            trie.insert(s);
            StringBuffer sb = new StringBuffer(s);

            reverseTrie.insert(String.valueOf(sb.reverse()));
        }
        ArrayList<Integer> array = new ArrayList<>();
        for (int i = 0; i < queries.length; i++) {
            int num = 0;
            if(queries[i].charAt(0) == '?'){
                StringBuffer sb = new StringBuffer(queries[i]);
                num = reverseTrie.search(sb.reverse().toString());
            }else {
                num = trie.search(queries[i]);
            }

            array.add(num);
            System.out.println(num);
        }
        answer = new int[array.size()];
        for (int i = 0; i < array.size(); i++) {
            answer[i] = array.get(i);
        }

        return answer;
    }

    public static void main(String[] args) {
        String[] words = {"frodo", "front", "frost", "frozen", "frame", "kakao"};
//        String[] queries = {"fro??", "????o", "fro???", "fr???", "pro?"};
        String[] queries = {"???", "fro??", "????o", "fr???", "fro???", "pro?"};
        solution(words, queries);
    }
}