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

문제

programmers.co.kr/learn/courses/30/lessons/60060

코딩테스트 연습 - 가사 검색

문제해결방법

처음에는 이 문제를 HashMap을 사용해서 풀려고 했지만 효율성 테스트에서 시간 초과, 메모리 초과가 발생해서 다른 방법으로 풀었다.

그래서 생각한 방법이 lower bound, upper bound 였다.

가사를 정렬하고 쿼리마다 정렬된 가사에서 lower bound와 upper bound를 계산해준 다음에 두 bound의 차이를 출력해줬다.

* 다 풀고나서 보니까 이 문제는 Trie로 푸는 문제인 것 같았다. 추후에 Trie로 푼 코드를 업데이트 해야겠다.

[21-02-08] Trie 코드 업데이트 완료

코드

package blogQueue;
 
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.List;
 
public class Solution_가사검색 {
 
    public static int compare(String s1, String s2) {
        int length = s2.length();
        if(s1.length()!=s2.length()) return s1.length() - s2.length();        
        if(s2.charAt(0)=='?' && s2.charAt(length-1)=='?') return 0;
        else {
            for(int i=0; i<s2.length(); i++) {
                if(s2.charAt(i)=='?') break;
                if(s1.charAt(i)!=s2.charAt(i)) return s1.charAt(i)-s2.charAt(i);
            }
        }
        return 0;
    }
    
    public static int[] solution(String[] words, String[] queries) {
        List<Integer> list = new ArrayList<>();
        String[] reverseWords = new String[words.length];
        for(int i=0; i<words.length; i++) reverseWords[i] = new StringBuilder(words[i]).reverse().toString();
        Arrays.sort(words, new Comparator<String>() {
            @Override
            public int compare(String arg0, String arg1) {
                if(arg0.length()!=arg1.length()) return Integer.compare(arg0.length(), arg1.length());
                return arg0.compareTo(arg1);
            }
        });
        Arrays.sort(reverseWords, new Comparator<String>() {
            @Override
            public int compare(String arg0, String arg1) {
                if(arg0.length()!=arg1.length()) return Integer.compare(arg0.length(), arg1.length());
                return arg0.compareTo(arg1);
            }
        });
        for(String query : queries) {
            int l = 0, r = words.length-1;
            int start = 0, end = 0;
            if(query.charAt(0)!='?') {
                while(l<=r) {
                    int mid = (l+r)>>1;
                    if(compare(words[mid], query)<0) l = mid+1;
                    else r = mid-1;
                }
                start = l;
                l = 0;
                r = words.length-1;
                while(l<=r) {
                    int mid = (l+r)>>1;
                    if(compare(words[mid], query)<=0) l = mid+1;
                    else r = mid-1;
                }
                end = l;
            } else {
                query = new StringBuilder(query).reverse().toString();
                while(l<=r) {
                    int mid = (l+r)>>1;
                    if(compare(reverseWords[mid], query)<0) l = mid+1;
                    else r = mid-1;
                }
                start = l;
                l = 0;
                r = words.length-1;
                while(l<=r) {
                    int mid = (l+r)>>1;
                    if(compare(reverseWords[mid], query)<=0) l = mid+1;
                    else r = mid-1;
                }
                end = l;
            }
            list.add(end-start);
        }
        int[] answer = new int[list.size()];
        for(int i=0; i<answer.length; i++) answer[i] = list.get(i);
        return answer;
    }
    
    public static void main(String[] args) {
        String[] words = {"frodo", "front", "frost", "frozen", "frame", "kakao"};
        String[] queries = {"fro??", "????o", "fr???", "fro???", "pro?"};
        System.out.println(Arrays.toString(solution(words, queries)));
 
    }
}

코드 - Trie

package blogQueue;
 
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
 
class TrieNode {
    int count;
    TrieNode[] nextNode;
    TrieNode() {
        this.count = 1;
        nextNode = new TrieNode[26];
    }
}
 
class Trie {
    int total;
    TrieNode root;
    Trie() {
        total = 0;
        root = new TrieNode();
    }
    public void insert(String s) {
        total++;
        TrieNode cur = root;
        for(int i=0; i<s.length(); i++) {
            int next = s.charAt(i)-'a';
            if(cur.nextNode[next]==null) cur.nextNode[next] = new TrieNode();
            else cur.nextNode[next].count++;
            cur = cur.nextNode[next];
        }
    }
    
    public int count(String s) {
        TrieNode cur = root;
        for(int i=0; i<s.length(); i++) {
            if(s.charAt(i)=='?') return cur.count;
            int next = s.charAt(i)-'a';
            if(cur.nextNode[next]==null) return 0;
            cur = cur.nextNode[next];
        }
        return cur.count;
    }
}
 
public class Solution_가사검색Trie {
 
    public static Trie[] trie, reverseTrie; 
 
    public static int[] solution(String[] words, String[] queries) {
        StringBuilder sb;
        trie = new Trie[100001];
        reverseTrie = new Trie[100001];
        List<Integer> list = new ArrayList<>();
        for(int i=0; i<trie.length; i++) {
            trie[i] = new Trie();
            reverseTrie[i] = new Trie();
        }
        for(String word : words) {
            sb = new StringBuilder(word);
            trie[word.length()].insert(word);
            reverseTrie[word.length()].insert(sb.reverse().toString());
        }
        for(String query : queries) {
            int length = query.length();
            if(query.charAt(0)=='?' && query.charAt(length-1)=='?') list.add(trie[length].total);
            else if(query.charAt(0)=='?') {
                sb = new StringBuilder(query);
                list.add(reverseTrie[length].count(sb.reverse().toString()));
            }
            else list.add(trie[length].count(query));
        }
        int[] answer = new int[list.size()];
        for(int i=0; i<answer.length; i++) answer[i] = list.get(i);
        return answer;
    }
    
    public static void main(String[] args) {
        String[] words = {"frodo", "front", "frost", "frozen", "frame", "kakao"};
        String[] queries = {"fro??", "????o", "fr???", "fro???", "pro?"};
        System.out.println(Arrays.toString(solution(words, queries)));
    }
}

코드에 대한 피드백이나 더 좋은 아이디어는 언제나 환영입니다.