[String, Trie] (프로그래머스 Lv 4) 가사 검색

문제

• 문제 링크
• 문제 설명: 가사에 사용된 모든 단어들이 담긴 배열 words와 찾고자 하는 키워드가 담긴 배열 queries가 주어질 때, 각 키워드 별로 매치된 단어가 몇 개인지 순서대로 배열에 담아 반환하는 함수 작성
• 입력
  • 가사 단어 제한 사항
    • words의 길이(가사 단어의 개수)는 2 이상 100,000 이하입니다.
    • 각 가사 단어의 길이는 1 이상 10,000 이하로 빈 문자열인 경우는 없습니다.
    • 전체 가사 단어 길이의 합은 2 이상 1,000,000 이하입니다.
    • 가사에 동일 단어가 여러 번 나올 경우 중복을 제거하고 words에는 하나로만 제공됩니다.
    • 각 가사 단어는 오직 알파벳 소문자로만 구성되어 있으며, 특수문자나 숫자는 포함하지 않는 것으로 가정합니다.
  • 검색 키워드 제한 사항
    • queries의 길이(검색 키워드 개수)는 2 이상 100,000 이하입니다.
    • 각 검색 키워드의 길이는 1 이상 10,000 이하로 빈 문자열인 경우는 없습니다.
    • 전체 검색 키워드 길이의 합은 2 이상 1,000,000 이하입니다.
    • 검색 키워드는 중복될 수도 있습니다.
    • 각 검색 키워드는 오직 알파벳 소문자와 와일드카드 문자인 '?' 로만 구성되어 있으며, 특수문자나 숫자는 포함하지 않는 것으로 가정합니다.
    • 검색 키워드는 와일드카드 문자인 '?'가 하나 이상 포함돼 있으며, '?'는 각 검색 키워드의 접두사 아니면 접미사 중 하나로만 주어집니다.
• 출력: 키워드별 매치된 단어 개수 배열

words	queries	result
["frodo", "front", "frost", "frozen", "frame", "kakao"]	["fro??", "????o", "fr???", "fro???", "pro?"]	[3, 2, 4, 1, 0]

풀이

• 전체 가사단어/검색키워드 길이의 합이 100만이 넘으므로 naive한 방법으로 탐색 시 시간 초과가 발생
• 트라이로 문자열을 탐색하되, 각 노드(문자1개)는 이후 남은 길이에 대해 단어의 개수를 map에 저장 (있는 그대로 탐색할 경우 시간초과)
• 와일드카드가 접두사/접미사에만 등장하므로 접두사 트라이와 접미사 트라이를 둘 다 구현해서 와일드 카드가 앞에 오면 접미사 트라이에서 탐색하고 와일드 카드가 뒤에 오면 접두사 트라이에서 탐색.
• 전처리할 때 검색 키워드 길이의 합 100만만큼 비용이 소모된다. 매칭 개수 찾는 것은 전처리보다 적게 듬
• 정답 코드

#include <cstring>
#include <string>
#include <vector>
#include <unordered_map>
#include <algorithm>
using namespace std;

struct TrieNode {
    bool terminal;
    TrieNode* children[26];
    unordered_map<int, int> m; // (len, cnt)
    
    TrieNode(): terminal(false) {
         memset(children, 0, sizeof(children));
    }
    
    void insert(int pos, const string& s) {
        if (pos < s.size()) {
            int idx = s[pos]-'a';
            if (children[idx] == 0)
                children[idx] = new TrieNode();
            children[idx]->insert(pos+1, s);
        } else terminal = true;
        m[s.size() - pos]++; // 남은 길이마다 단어 수 증가
    }
    
    int search(int pos, const string& s) {
        if (s[pos] == '?') return m[s.size() - pos]; // 다음 노드에 대해 ?이면 단어 수 반환
        else if (children[s[pos]-'a'])
            return children[s[pos]-'a']->search(pos+1, s);
        return 0;
    }
};

vector<int> solution(vector<string> words, vector<string> queries) {
    TrieNode* prefixTrie = new TrieNode();
    TrieNode* suffixTrie = new TrieNode();
    for (auto& w : words) {
        prefixTrie->insert(0, w);
        reverse(w.begin(), w.end());
        suffixTrie->insert(0, w);
    }
    vector<int> answer;
    for (auto& q : queries) {
        if (q[0] == '?') { // 접두사에 와일드카드가 있으면 접미사 트라이에서 탐색
            reverse(q.begin(), q.end());
            answer.push_back(suffixTrie->search(0, q));
        } else { // 접미사에 와일드카드가 있으면 접두사 트라이에서 탐색
            answer.push_back(prefixTrie->search(0, q));
        }
    }
    return answer;
}