카카오 신입 공채 3차 코딩 테스트 풀이 - 5번문제

카카오 3차 코딩 테스트에서는 트라이라는 자료구조를 알고 있는지에 대해 거의 직접적으로 물어보는 수준의 문제를 출제하였습니다. 트라이는 빈 문자열을 루트로 가지고, 각 단어 1글자당 1개의 depth를 가지는 트리구조입니다. 트라이는 바이너리 서치 트리를 문자열에 적용할 경우 문자열의 길이  $M$ 만큼의 시간복잡도가 곱연산으로 늘어나는 대체 자료구조로 사용할 수 있으며, 검색의 시간복잡도는 $O(M)$ 입니다.

(출처 : 위키피디아)

트라이의 구현은 링크드리스트로 해도 되고, 이 문제의 경우에는 해쉬를 이용해도 알파벳이 26글자 밖에 없으므로 전혀 지장이 없습니다. 문자열의 맨 끝에 검색이 끝났다는 표시 심볼(저 같은 경우는 $를 사용하였습니다.) 를 추가하여 별도의 처리를 해줘도 됩니다. 자동 완성이 되는 시점은

1) 자식 노드가 1개이며, 그 문자열이 $인 경우 (문자열 전체를 다 입력했을 때)

2) 자식 노드가 1개이며, 그 문자열이 추가된 횟수가 1인 경우(이 문자열 이후로는 오직 1번만 추가되었음을 의미)

입니다. 이 2가지를 캐치하여 각각의 단어에 대해 count를 해주면 정답이 나옵니다.

class Trie
{
public:
    set<char> words;
    vector<int> cnt;

    Trie() { cnt.resize(26); }
};
long long solve(vector<string> words)
{
    map<int, map<char, Trie> > trie;

    for (auto word : words)
    {
        word += '$';
        int st = 0;
        for (int i=0; i<word.length() -1; i++)
        {
            trie[st][word[i]].words.insert(word[i + 1]);
            if (word[i + 1] == '$') break;
            trie[st][word[i]].cnt[word[i + 1] - 'a']++;
            st++;
        }
    }
    long long ans = 0;
    for (auto word : words)
    {
        long long sum = 0;
        int st = 0;

        for (int i = 0; i < word.length(); i++)
        {
            sum++;
            if (trie[st][word[i]].words.size() == 1)
            {
                char spell = *trie[st][word[i]].words.begin();
                if (spell=='$' || trie[st][word[i]].cnt[spell - 'a'] == 1) break;
            }
            st++;
        }
        ans += sum;
    }
    return ans;
}