2020 카카오 신입사원 코딩테스트 - 가사 검색

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

코딩테스트 연습 - 가사 검색 | 프로그래머스

 

이번에는 카카오 코딩 테스트 4번째 문제 풀이 입니다.

 

이 문제는 'Trie' 자료구조를 사용하는 문제입니다.

 

문자열 알고리즘은 크게 KMP, Trie, Rabin-Karp, Aho-Corasick, Suffix Array 가 있습니다.

Trie를 사용하는 문제는 알고리즘 대회나 코딩 테스트에서 흔히 볼 수 있는 문제는 아닙니다.

 

카카오 코테에서 이 문제는 두 번이나 나온걸 보면 좀 특이하긴 합니다.

 

Trie 는 여러 문자열을 빠르게 찾을 수 있도록 저장하고 있는 자료구조 입니다.

어떤 문자열이 있나 찾을 때 아주 빠르게 찾을 수 있습니다.

 

이 알고리즘에 대한 자세한 설명은 곧 올릴 예정입니다. 

 

여기서는 아래 그림을 보면서 간단하게 설명해보겠습니다.

 

 

 

출처 : https://cgi.cse.unsw.edu.au/~cs9024/19t0/probs/prob07/Pics/trie.png

 

루트에서 시작하여 노드에 저장되어 있는 알파벳을 빨간색 노드를 만날 때까지 하나씩 이어붙여 보면

'a', 'air', 'an', 'anon', 'car', 'ma', 'man', 'mane', 'oh', 'or', 'orate' 문자열들을 만들 수 있습니다.

이렇게 문자열을 저장해 놓는 것입니다.

 

이때, 만약 'air' 또는 'ail' 등의 문자열이 있는가라는 쿼리가 들어오면 아주 쉽게 찾을 수 있죠!

 

아무리 많은 문자열이 저장되어 있더라도, 이 자료구조를 사용하면, 

시간복잡도 O(N) 만에 판별할 수 있습니다. 여기서 N은 찾고자 하는 문자열의 길이 입니다.

 

 

이제 문제를 한 번 보겠습니다.

이 문제는 쿼리가 좀 특이하게 들어옵니다.

접두 또는 접미에 '?'를 달고 오는데, '?'는 하나의 임의의 문자와 매칭됩니다.

 

접두에 달고올 때, 해결하기가 쉽지 않습니다.

이렇게 되면 트라이 자료구조를 사용하는 의미가 없어질 정도로 찾는 시간이 늘어나게 됩니다!

 

이 문제는 들어온 단어 그대로 트라이(forward)에 넣어주고, 거꾸로 뒤집은 단어를 또다른 트라이(backward)에 넣어줍니다.

즉, 트라이 자료구조를 두 개를 운영해서 해결하였습니다.

 

접두사에 '?'를 달고 들어오는 쿼리는 backward 트라이에서 찾고,

접미사에 '?'가 달려있으면 forward 트라이에서 찾기를 시도하면 됩니다!

 

여기서 또 하나의 문제가 생깁니다.

'?'가 시작하는 시점에서, '?'의 길이와 일치하는 문자열의 갯수를 구해야 합니다.

물론 트리 순회를 통해 쉽게 알 수 있지만, 이를 시간 초과 발생하게 하는 데이터를 충분히 만들 수 있습니다.

 

이 문제를 해결하기 위한 많은 방법이 있지만, 여기서는 각 노드에 unordered_map을 두어 해결했습니다.

이 변수는 문자열 길이별로 문자열의 갯수를 저장하고 있습니다.

 

그렇게 되면, 트리 순회를 할 필요가 없습니다.

만약, 'abc???'를 찾는다고 하죠. 이는 길이가 6인 문자열 입니다.

'?'를 만난 시점에 노드의 mp[6]을 찾아가면 총 길이 6이면서 접두사 'abc'를 갖는 문자열의 총 갯수를 알 수 있습니다.

 

 

#include <string>
#include <vector>
#include <algorithm>
#include <unordered_map>
 
using namespace std;
 
const int go_max = 26;
 
struct trie {
    bool output;
    trie* go[go_max];
    unordered_map<int,int> mp;
 
    trie() {
        fill(go, go + go_max, nullptr);
        output = false;
    }
 
    ~trie() {
        for(int i = 0; i < go_max; i++)
            if(go[i]) delete go[i];
    }
 
    void insert(const char *key, int l) {
        if(*key == '\0') {
            output = true;
            return;
        }
        int idx = *key - 'a';
        mp[l] += 1;
        if(!go[idx]) {
            go[idx] = new trie;
        }
        go[idx]->insert(key + 1, l);
    }
 
    int find(const char *key, int l) {
        if(*key == '\0') return output;
        int ret = 0;
        int idx = *key - 'a';
        if(*key == '?') {
            ret += mp[l];
        }
        else {
            if(!go[idx]) return 0;
            ret = go[idx]->find(key + 1, l);
        }
        return ret;
    }
};
 
vector<int> solution(vector<string> w, vector<string> q) {
    trie* root_f = new trie;
    trie* root_r = new trie;
 
    int wl = w.size(), ql = q.size();
    for(int i = 0; i < wl; i++) {
        string s = w[i];
        root_f->insert(&s[0], s.size());
        reverse(s.begin(), s.end());
        root_r->insert(&s[0], s.size());
    }
 
    vector<int> ans(ql);
    for(int i = 0; i < ql; i++){
        string query = q[i];
        if(query[0] != '?'){
            ans[i] = root_f->find(&query[0], query.size());
        }
        else {
            reverse(query.begin(), query.end());
            ans[i] = root_r->find(&query[0], query.size());
        }
    }
    return ans;
}