[String, Trie] (프로그래머스 Lv 4) 자동완성

문제

• 문제 링크
• 문제 설명: 단어 목록 words가 주어질 때, 매 단어를 자동완성하기 위해 입력해야 하는 글자 수의 총합을 반환하는 함수를 작성하라. 접두사가 겹치는 경우는 겹치지 않을 때까지 입력하거나 단어를 모두 입력해야 한다.
• 입력
  • 학습과 검색에 사용될 중복 없는 단어 N개가 주어진다. 
  • 모든 단어는 알파벳 소문자로 구성되며 단어의 수 N과 단어들의 길이의 총합 L의 범위는 다음과 같다.
  • 2 <= N <= 100,000
  • 2 <= L <= 1,000,000
• 출력: 단어를 자동완성하기 위해 입력해야하는 글자 수의 총합

words	result
["go","gone","guild"]	7
["abc","def","ghi","jklm"]	4
["word","war","warrior","world"]	15

풀이

• 단어들의 길이 총합이 100만이므로 O(N)으로 문제를 해결해야 함
• 트라이로 구현해서 단어들이 겹치는 문자들에 대해 카운트 변수를 유지(=트라이노드의 멤버변수)
  • 삽입할 때마다 카운트 변수를 증가
• 검색할 때 단어의 끝 또는 카운트가 1이 되는 순간(=후보 단어가 하나 밖에 없는 경우) 길이를 계산해서 반환하도록 한다.
• 시간 복잡도: O(L)
• 정답 코드

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

struct TrieNode {
    bool terminal;
    TrieNode* children[26];
    int count;
    
    TrieNode() : terminal(false), count(0) {
        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;
        count++; // 단어 카운트
    }
    
    int search(int pos, const string& s) {
        // base case: 찾는 단어의 끝이거나 후보 단어가 1개 밖에 없으면 재귀 종료
        if ((pos == s.size() and terminal) or count == 1) return 0;
        return children[s[pos]-'a']->search(pos + 1, s) + 1;
    }
};

int solution(vector<string> words) {
    TrieNode* root = new TrieNode();
    for (auto& s : words) root->insert(0, s);
    int answer = 0;
    for (auto& s : words) answer += root->search(0, s);
    return answer;
}