[자바] 프로그래머스 - 가사 검색 [코딩테스트 연습 Lv4]

문제 : Programmers-가사검색

 

 

  트라이(Trie) 개념으로 구현하면 된다. 내 경우 words에 들어있는 String으로 트라이 트리를 구성했는데, 각 길이마다 별도로 앞에서 부터 진행한 트리와 뒤에서 부터 진행한 트리를 둘 다 구성했다. 만약 '*' 같은것도 있으면 길이마다 구성한게 의미가 없었을테지만, 이 문제의 경우 queries의 길이가 고정되어 있으므로 길이마다 별도로 구성하는 것이 더 빠르게 연산할 수 있다. 또한 '?'가 접미사 혹은 접두사에 있으므로 둘 다 찾기 위해서는 앞에서 부터 구성한 트라이와 뒤에서 부터 구성한 트라이 둘 다 유지해야 한다. 접미사에 '?'가 있다면 뒤에서 부터 구성한 트라이 트리를 봐야 하고, 접두사에 있다면 앞에서 부터 구성한 트라이 트리를 봐야 한다. 추가로 cnt를 넣어 해당 노드 이하에 있는 모든 개수를 체크하도록 했다. 그래서 더 밑으로 내려가지 않아도, '?'가 나오는 위치에서 cnt만 읽으면 된다.

 

  예를들어 words가 ["abcd", "abcc", "abdd", "ab", "ac", abcde"] 인 경우 다음과 같이 트라이 트리를 유지한다. 길이별로 나누고, 거기서 앞부분 부터 유지한 트라이(녹색)과, 뒤에서 부터 유지한 트라이(파란색)로 나뉜다. 각 노드의 숫자는 해당 루트로 진행 시 리프노드가 몇 개 인지를 나타낸다.

  예를들어 queries에 "abc?"가 존재했다면 다음과 같이 찾으면 된다. 우선 길이가 4이므로 길이4에 대해 유지한 트라이로 가고, 접미사이므로 앞부터 유지한 트라이로 간다. 그 후 a->b->c를 진행하고 그 다음 '?' 이므로 더이상 진행해볼 필요 없이 cnt를 보면 이하 리프노드가 2개임을 알 수 있다.

 

 

코드 : github

/**
 * 문제 출처: 프로그래머스 코딩 테스트 연습, https://programmers.co.kr/learn/challenges
 */
class Node {
    int cnt;
    Node[] child;
    public Node() {
        child = new Node['z'-'a'+1];
        cnt = 0;
    }
    public Node makeAndReturnChild(char c) {
        cnt++;
        if (child[c-'a'] == null) {
            child[c-'a'] = new Node();
        }
        return child[c-'a'];
    }
    public Node getChild(char c) {
        return child[c-'a'];
    }
}

class StringChecker {
    Node strChk;
    Node reverseStrChk;
    public StringChecker() {
        strChk = new Node();
        reverseStrChk = new Node();
    }

}

class Solution {
    private static final int MAX_WORDS_LENGTH = 10000;
    private StringChecker[] checkers = new StringChecker[10001];

    private void init() {
        checkers = new StringChecker[MAX_WORDS_LENGTH+1];
        for (int i = 1; i <= MAX_WORDS_LENGTH; i++) {
            checkers[i] = new StringChecker();
        }
    }

    private void insert(String s) {
        Node pt = checkers[s.length()].strChk;
        Node reversePt = checkers[s.length()].reverseStrChk;
        for (int i = 0; i < s.length(); i++) {
            pt = pt.makeAndReturnChild(s.charAt(i));
            reversePt = reversePt.makeAndReturnChild(s.charAt(s.length()-1-i));
        }
    }

    private int getCnt(String s) {
        Node pt = checkers[s.length()].strChk;
        for (int i = 0; i < s.length(); i++) {
            char c = s.charAt(i);
            if (c == '?')
                return pt.cnt;
            Node tmp = pt.getChild(c);
            if (tmp == null)
                break;
            pt = tmp;
        }
        return 0;
    }

    private int getCntFromReverse(String s) {
        Node reversePt = checkers[s.length()].reverseStrChk;
        for (int i = 0; i < s.length(); i++) {
            char c = s.charAt(s.length()-1-i);
            if (c == '?')
                return reversePt.cnt;
            Node tmp = reversePt.getChild(c);
            if (tmp == null)
                break;
            reversePt = tmp;
        }
        return 0;
    }

    public int[] solution(String[] words, String[] queries) {
        init();
        for (int i = 0; i < words.length; i++) {
            insert(words[i]);
        }
        int[] answer = new int[queries.length];
        for (int i = 0; i < queries.length; i++) {
            answer[i] = queries[i].charAt(0)=='?' ? getCntFromReverse(queries[i]) : getCnt(queries[i]);
        }
        return answer;
    }
}