가사 검색 (프로그래머스) JAVA

https://school.programmers.co.kr/learn/courses/30/lessons/60060

프로그래머스

 

문제 설명

이 문제는 words가 주어지면 이 words에 맞는 queries가 몇 개 있는지 찾는 문제입니다.

fro??은 frodo, front, frost 이렇게 3개가 존재합니다.

 

문제에 대한 아이디어

제일 처음에는 queries 하나 당 words를 돌면서 똑같은 것이 있는지 체크를 할려고 했습니다. 하지만 이 문제는 효율성 테스트도 있는 문제입니다. 단어 개수를 W, 단어 길이를 L, 쿼리 개수를 Q라고 하면 O(WLQ)가 됩니다. 이 떼, 전체 단어 길이 합은 최대 1000000 이므로 WL은 1000000이고 Q는 최대 100000이므로 시간 제한을 통과하지 못합니다.

어떻게 풀지 고민하다가 책을 찾아봤습니다. 프로그래머스 코딩 테스트 문제 풀이 전략 이라는 책에서 이 문제는 트라이를 풀면 쉽게 풀 수 있다고 했습니다. 트라이는 들어만 봤지 처음 써봐서 이 기회에 공부해 봤습니다.

 

트라이란?

트라이(Trie)는 검색을 목적으로하는 트리(Tree) 자료 구조 중 하나입니다. 트라이는 문자열 검색에 사용되며, 문자열을 이진 검색 트리(Binary Search Tree)로 표현하면 검색 및 삽입 시간이 문자열 길이에 관계없이 상수 시간( O(1))으로 유지되는 장점이 있습니다.

위의 예시를 트라이 자료구조로 변경하면 아래의 그림과 같은 형식 입니다.

각 노드는 하나의 char입니다.

fr 까지 일치하는 것은 총 5개가 있는 것을 알 수 있습니다. 이렇게 트라이 자료구조는 일치하는 단어를 쉽게 찾을 수 있습니다. 

맨 위 루트에서 보면 ?가 5개일 경우 5개의 단어를 만들 수 있고 ?가 6개일 경우 1개의 단어를 만들 수 있습니다. 이 부분을 Map으로 저장해노면 지금 이 노드에서 몇개의 단어를 만들 수 있는지 쉽게 찾을 수 있습니다.

private final Map<Integer, Integer> terminals = new HashMap<>();  //? 개수별로 완성되는 단어의 수
private final Map<Character, Node> children = new HashMap<>();   // 자신의 자식

• add 함수

public void add(String word, int offset){
       int length = word.length() - offset; // 단어에 필요한 ? 갯수
       terminals.put(length, terminals.getOrDefault(length, 0)+1); 
       if(length > 0){ // ? 갯수가 필요하면
           char c = word.charAt(offset); 
           Node child = children.getOrDefault(c, new Node()); // 자식 있는지 검사
           child.add(word, offset+1); 
           children.put(c, child); // 자식 추가
       }
}

word와 단어 내 문자 인덱스 offset을 입력받아 트라이를 구성 합니다. 트라이의 한 노드는 각 인덱스의 문자입니다.

이렇게 add를 하면서 트라이 자료구조를 생성할 수 있습니다!!

 

전체 코드

import java.util.*;

class Solution {
    private static class Node {
        private final Map<Integer, Integer> terminals = new HashMap<>();
        private final Map<Character, Node> children = new HashMap<>();
        
        //단어와 단어 내 문자 인덱스
        public void add(String word, int offset){
            int length = word.length() - offset;
            terminals.put(length, terminals.getOrDefault(length, 0)+1);
            if(length > 0){
                char c = word.charAt(offset);
                Node child = children.getOrDefault(c, new Node());
                child.add(word, offset+1);
                children.put(c, child);
            }
        }
        
        // 완성하는 단어가 몇개인지 쓰는 메소드
        public int count(String query, int offset){
            if(query.charAt(offset) == '?'){
                return terminals.getOrDefault(query.length() - offset, 0);
            }
            char c = query.charAt(offset);
            if(!children.containsKey(c)) return 0;
            return children.get(c).count(query, offset+1);
        }
    }
    
    // 처음부터 ?가 나오는 경우 처리
    private int count(String query, Node trie, Node reversedTrie) {
        if (query.startsWith("?")) {
            return reversedTrie.count(new StringBuilder(query).reverse().toString(), 0);
        }
        return trie.count(query, 0);
    }
    
    public int[] solution(String[] words, String[] queries) {
        Node trie = new Node();
        Node reversedTrie = new Node();
        for(String word : words){
            trie.add(word, 0);
            reversedTrie.add(new StringBuilder(word).reverse().toString(), 0);
        }
        
        return Arrays.stream(queries)
                 .mapToInt(query -> count(query, trie, reversedTrie))
                 .toArray();
    }
}