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

문제 링크

#include <string>
#include <vector>
#include <algorithm>

using namespace std;

class Trie {
private:
    Trie* child[26];  // 자식 노드 주소 배열 (배열로 바로 접근 가능, 인덱스는 소문자 알파벳 순서에 대응)
    int count;  // 첫 문자부터 해당 노드의 문자까지를 접두사로 가지는 문자들의 개수 (= ) 
public:
    Trie() : child(), count(0) {} // child 는 null 로, count 는 0 으로 초기화
    /* 트라이 만들기 */
    void Insert(string str) {
        Trie* now = this;
        for (char ch : str) {
            now->count++; // 방문할 때마다 1 증가해주면 됨
            if (now->child[ch - 'a'] == nullptr)
                now->child[ch - 'a'] = new Trie(); // 없다면 새로 생성
            now = now->child[ch - 'a']; // 트리 타고 내려감
        }
    }

    int Search(string str) {
        Trie* now = this;
        for (char ch : str) {
            if (ch == '?')  // ? 면 검사하려는 접두사가 끝났다는 것
                return now->count; // 현재 방문중인 노드의 count 를 리턴하고 끝내면 된다.
            now = now->child[ch - 'a'];
            if (now == nullptr) // 트라이의 끝까지 도달 할 때까지도 해당 접두사와 일치하는 것을 만나지 못했다면 해당 접두사로 시작하는 단어는 없는 것
                return 0;
        }
        return -1; // 코드가 정상적이라면 여기에 걸릴 일은 없다.
    }
};

Trie TrieRoot[10000]; // 접두사 찾는 길이별 트리들 (문제에서 최대 길이 10000 이라고 제한함) 
Trie ReTrieRoot[10000]; // 접미사 찾는 길이별 트리들! 거꾸로 뒤집은 문자열들을 넣어 접두사 찾듯이 찾을 것
vector<int> solution(vector<string> words, vector<string> queries) {
    vector<int> answer;
    /* 트라이 트리 만들기 (개수 구해놓기) */
    for (string str : words) { 
        TrieRoot[str.length() - 1].Insert(str); // 길이에 맞는 트리에 원래 순서대로 추가

        reverse(str.begin(), str.end()); // 문자열 뒤집기
        ReTrieRoot[str.length() - 1].Insert(str); // 길이에 맞는 트리에 역순서으로 추가
    }
    /* 트라이 트리 탐색 */
    for (string query : queries) {
        if (query[0] != '?') // 접두사 : TrieRoot[str.length() - 1] 트리 탐색
            answer.push_back(TrieRoot[query.length() - 1].Search(query));
        else {  // 접미사 : ReTrieRoot[str.length() - 1] 트리 탐색
            reverse(query.begin(), query.end()); // 뒤집기. ??and 가 dna?? 가 되도록
            answer.push_back(ReTrieRoot[query.length() - 1].Search(query));
        }
    }
    return answer;
}