프로그래머스 : 가사 검색

https://programmers.co.kr/learn/courses/30/lessons/60060?language=cpp

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

 

'?'는 query 문자열에 무조건 들어가고, 중간에는 들어가지않는 조건을 염두에 두고 풀었습니다.

효율성 테스트 때문에 Trie를 최대 길이만큼 선언해준 뒤, 각 word의 길이에 해당하는 Trie에만 삽입해주었습니다.

'?'는 query 문자열의 뒤에서 시작할 수도 있기 때문에, 문자열을 역순으로 뒤집은 것도 별도의 Trie에 삽입해주었습니다.

예를 들어, 'abc'는 역순으로 뒤집은 'cba'도 역순 Trie에 삽입하였습니다.

삽입을 하며, 지나가는 경로의 cnt를 1씩 계속 증가시켜주었습니다.

그렇게 되면, 루트에는 현재 Trie에 들어간 모든 단어의 개수가 나옵니다.

검색을 할 때는, query 문자열에서 '?'를 만났을 때의 cnt를 반환해주면 됩니다. 어차피 길이별로 구분해서 넣어주기 때문입니다.

 

 

#include <string>
#include <vector>
using namespace std;

class Trie {
private:
    int cnt = 0;
    Trie *alp[26];
public:
    ~Trie() {
        for(int i=0; i<26; i++) {
            if(alp[i] != NULL) delete alp[i];
        }
    }
    void insert(string& word, int idx, bool reverse) {
        cnt++; // 지나간 경로 ++; 루트의 cnt에는 모든 단어의 개수가 들어감
        if(idx == (reverse ? -1 : word.size())) return;
        int alp_idx = word[idx] - 'a';
        if(alp[alp_idx] == NULL) alp[alp_idx] = new Trie();
        alp[alp_idx]->insert(word, reverse ? idx - 1 : idx + 1, reverse);
    }
    
    int count(string& query, int idx, bool reverse) {
        if(query[idx] == '?') return cnt; // '?'를 만났다면 현재지점의 cnt 반환
        else if(alp[query[idx]-'a'] != NULL)
            return alp[query[idx]-'a']->count(query, reverse ? idx - 1 : idx + 1, reverse);
        else return 0;
    }
};

vector<int> solution(vector<string> words, vector<string> queries) {
    vector<int> answer;
    
    Trie* trie[10000]; // 각 단어의 길이만큼 생성
    Trie* rtrie[10000];
    for(int i=0; i<words.size(); i++) {
        int idx = words[i].size()-1;
        if(trie[idx] == NULL) {
            trie[idx] = new Trie();
            rtrie[idx] = new Trie();
        }
        trie[idx]->insert(words[i], 0, false);
        rtrie[idx]->insert(words[i], words[i].size()-1, true);
    }
    for(int i=0; i<queries.size(); i++) {
        int idx = queries[i].size()-1;
        int ret = 0;
        if(queries[i][0] == '?' && queries[i][idx] == '?') { // '?'로만 이루어진 문자열
            if(trie[idx] != NULL)
                ret = trie[idx]->count(queries[i], 0, false);
        } else {
           if(trie[idx] != NULL)
                ret = queries[i][0] == '?' ? rtrie[idx]->count(queries[i], idx, true) : trie[idx]->count(queries[i], 0, false);
        }
        answer.push_back(ret);
    }
    return answer;
}