가사검색(Trie tree)

문제 설명

이 문제는 트라이 트리를 이용한 문제이다.
숫자 탐색에 용이한 이진트리와 달리, 트라이는 Character로 트리구조로 만들어, 단어 탐색에 용이하다. 이 문제에서는 단어의 길이 마다, 트라이 구조를 만들어서 해결하였다. 트리에서 단어의 문자들을 거쳐갈 때 마다, count하여 문자를 포함한 단어가 몇개인지 알 수 있도록 설계하였다.

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

문제 설명
[본 문제는 정확성과 효율성 테스트 각각 점수가 있는 문제입니다.]

친구들로부터 천재 프로그래머로 불리는 "프로도"는 음악을 하는 친구로부터 자신이 좋아하는 노래 가사에 사용된 단어들 중에 특정 키워드가 몇 개 포함되어 있는지 궁금하니 프로그램으로 개발해 달라는 제안을 받았습니다.
그 제안 사항 중, 키워드는 와일드카드 문자중 하나인 '?'가 포함된 패턴 형태의 문자열을 뜻합니다. 와일드카드 문자인 '?'는 글자 하나를 의미하며, 어떤 문자에도 매치된다고 가정합니다. 예를 들어 "fro??"는 "frodo", "front", "frost" 등에 매치되지만 "frame", "frozen"에는 매치되지 않습니다.

가사에 사용된 모든 단어들이 담긴 배열 words와 찾고자 하는 키워드가 담긴 배열 queries가 주어질 때, 각 키워드 별로 매치된 단어가 몇 개인지 순서대로 배열에 담아 반환하도록 solution 함수를 완성해 주세요.

가사 단어 제한사항
words의 길이(가사 단어의 개수)는 2 이상 100,000 이하입니다.
각 가사 단어의 길이는 1 이상 10,000 이하로 빈 문자열인 경우는 없습니다.
전체 가사 단어 길이의 합은 2 이상 1,000,000 이하입니다.
가사에 동일 단어가 여러 번 나올 경우 중복을 제거하고 words에는 하나로만 제공됩니다.
각 가사 단어는 오직 알파벳 소문자로만 구성되어 있으며, 특수문자나 숫자는 포함하지 않는 것으로 가정합니다.
검색 키워드 제한사항
queries의 길이(검색 키워드 개수)는 2 이상 100,000 이하입니다.
각 검색 키워드의 길이는 1 이상 10,000 이하로 빈 문자열인 경우는 없습니다.
전체 검색 키워드 길이의 합은 2 이상 1,000,000 이하입니다.
검색 키워드는 중복될 수도 있습니다.
각 검색 키워드는 오직 알파벳 소문자와 와일드카드 문자인 '?' 로만 구성되어 있으며, 특수문자나 숫자는 포함하지 않는 것으로 가정합니다.
검색 키워드는 와일드카드 문자인 '?'가 하나 이상 포함돼 있으며, '?'는 각 검색 키워드의 접두사 아니면 접미사 중 하나로만 주어집니다.
예를 들어 "??odo", "fro??", "?????"는 가능한 키워드입니다.
반면에 "frodo"('?'가 없음), "fr?do"('?'가 중간에 있음), "?ro??"('?'가 양쪽에 있음)는 불가능한 키워드입니다.

입출력 예시

입력 :
words = ["frodo", "front", "frost", "frozen", "frame", "kakao"]
queries = ["fro??", "????o", "fr???", "fro???", "pro?"]

출력 :
result = [3, 2, 4, 1, 0]

구현

Java

import java.util.Arrays;

class TrieNode {
    TrieNode[] child = new TrieNode[26];  // 알파벳 26개에 해당하는 자식노드를 가질 수 있다.
    int count = 0;  // 해당 노드를 거쳐갈 수 있는 단어의 개수를 저장한다.
}

public class TriesTree {
    static TrieNode[] root = new TrieNode[10001];  // 단어의 길이가 1~10000까지이므로 길이에 따른 루트노드를 가진 배열을 생성한다.
    static TrieNode[] reverseRoot = new TrieNode[10001];  // 역방향 검색을 위한 Trie도 생성한다.

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

        buildTries(words);  // 주어진 단어들로 Trie를 만든다.

        int[] answer = new int[queries.length];
        int ansIdx = 0;
        for (String query : queries) {
            if (query.charAt(0) != '?') {  // prefix 검색의 경우
                answer[ansIdx++] = search(root[query.length()], query);
            } else {  // suffix 검색의 경우
                String reversedQuery = new StringBuilder(query).reverse().toString();  // 역순으로 문자열을 만든다.
                answer[ansIdx++] = search(reverseRoot[query.length()], reversedQuery);
            }
        }
        System.out.println(Arrays.toString(answer));  // 검색한 결과를 출력한다.
    }
    
	// Trie 구조로 정리
    private static void buildTries(String[] words) {
        for (String word : words) {
            int len = word.length();
            if (root[len] == null) {  // 단어 길이에 해당하는 루트노드가 없으면 생성한다.
                root[len] = new TrieNode();
                reverseRoot[len] = new TrieNode();  // 역방향 검색을 위한 Trie도 생성한다.
            }

            insertTrie(root[len], word);  // 단어를 Trie에 추가한다.
            insertTrie(reverseRoot[len], new StringBuilder(word).reverse().toString());  // 역순으로도 Trie에 추가한다.
        }
    }

    //삽입 함수
    private static void insertTrie(TrieNode node, String word) {
        for (char ch : word.toCharArray()) {
            node.count++;  // 해당 노드를 거쳐갈 수 있는 단어의 개수를 1 증가시킨다.
            int idx = ch - 'a';
            if (node.child[idx] == null) {  // 해당 자식노드가 없으면 새로 생성한다.
                node.child[idx] = new TrieNode();
            }
            node = node.child[idx];  // 다음 자식노드로 이동한다.
        }
        node.count++;  // 마지막 노드에서도 count를 증가시켜준다.
    }

    //검색 함수
    private static int search(TrieNode node, String query) {
        for (char ch : query.toCharArray()) {
            if (node == null) { //찾는 노드가 null이면 0반환
                return 0;
            }
            if (ch == '?') {
                return node.count; // ?를 만나면 count값 반환
            }
            int idx = ch - 'a';
            node = node.child[idx]; // 다음 자식노드로 이동한다.
        }
        return node == null ? 0 : node.count; // 코드가 여기까지 올 수 없다. 반환값은 위와 같음
    }
}

Python

class TrieNode:
    def __init__(self):
        self.child = [None] * 26
        self.count = 0

root = [None] * 10001
reverseRoot = [None] * 10001

def main():
    words = ["frodo", "front", "frost", "frozen", "frame", "kakao"]
    queries = ["fro??", "????o", "fr???", "fro???", "pro?"]
    
    buildTries(words)

    answer = [0] * len(queries)
    ansIdx = 0
    for query in queries:
        if query[0] != '?':
            answer[ansIdx] = search(root[len(query)], query)
        else:
            reversedQuery = query[::-1]
            answer[ansIdx] = search(reverseRoot[len(query)], reversedQuery)
        ansIdx += 1
    print(answer)

def buildTries(words):
    for word in words:
        length = len(word)
        if not root[length]:
            root[length] = TrieNode()
            reverseRoot[length] = TrieNode()

        insertTrie(root[length], word)
        insertTrie(reverseRoot[length], word[::-1])

def insertTrie(node, word):
    for ch in word:
        node.count += 1
        idx = ord(ch) - ord('a')
        if not node.child[idx]:
            node.child[idx] = TrieNode()
        node = node.child[idx]
    node.count += 1

def search(node, query):
    for ch in query:
        if not node:
            return 0
        if ch == '?':
            return node.count
        idx = ord(ch) - ord('a')
        node = node.child[idx]
    return node.count if node else 0

if __name__ == '__main__':
    main()