삼성 소프트웨어 (SW) 역량테스트 B형 준비! ② (Trie)

문제

https://www.acmicpc.net/problem/5052

 

풀이

N = 10^5 이기 때문에, O(NlogN) 이하의 알고리즘이 필요합니다.

Trie 라는 자료구조는 문자열의 검색 및 업데이트에 최적화 되어있습니다.

 

 

Trie란? 문자열의 Search Tree 입니다. 위키피디아 참조 위 그림처럼, 단어들의 첫 글자부터 하나씩 depth가 깊어지며 node를 만들기 시작합니다. Trie의 Update 및 Search 등은 Trie의 트리의 높이만큼의 시간뿐이 걸리지 않습니다. (매우빠름)

 

위 문제에서는, 각 번호들로 Trie 를 만들면서 단어가 끝나는 지점에 표시를 해줍니다.

void Node::insert(char* word) {
    if(word[0] == '\0') {
        this->isLast = true;
        return ;
    }
    
    char ch = word[0];
    int ch_int = ch - '0';
    
    if ( this->child[ch_int] != NULL ) {
        this->child[ch_int]->insert(word + 1);
    } else {
        this->child[ch_int] = new Node();
        this->child[ch_int]->value = ch; 
        this->child[ch_int]->insert(word + 1);
    }
}

모든 번호들을 Insert 한 후에, 다시 각 단어별로 탐색을 시작합니다.

이 때 탐색하는 도중 번호의 끝이 아님에도 불구하고, 다른 단어의 끝나는 지점에 도착했다면 다른 번호가 탐색하는 번호의 접두어를 이루고 있다는 뜻이 됩니다.

bool Node::validate(char* word) {
    if(word[0] == '\0') return true; // 탐색하는 단어의 끝인가?
    if(this->isLast) return false; // 탐색하는 도중, isLast=true 인 node를 만났다. 즉 다른단어의 끝나는 지점을 만났다!
    int ch_int = word[0] - '0';
    return this->child[ch_int]->validate(word+1);
}

 

 

코드

/*
https://www.acmicpc.net/problem/5052

N = 10^5 이기 때문에,
O(NlogN) 이하의 알고리즘이 필요하다.

String의 접두어 관련 문제는 Trie 로 해결이 가능할 것으로 보인다.
*/

#include <iostream>
#include <fstream>

using namespace std;

ifstream fin;
ofstream fout;

class Node {
    char value;
    bool isLast = false;
    Node *child[26]; 
    public:
        void insert(char* word);   
        bool validate(char* word);
};

void Node::insert(char* word) {
    if(word[0] == '\0') {
        this->isLast = true;
        return ;
    }
    
    char ch = word[0];
    int ch_int = ch - '0';
    
    if ( this->child[ch_int] != NULL ) {
        this->child[ch_int]->insert(word + 1);
    } else {
        this->child[ch_int] = new Node();
        this->child[ch_int]->value = ch; 
        this->child[ch_int]->insert(word + 1);
    }
}

bool Node::validate(char* word) {
    if(word[0] == '\0') return true;
    if(this->isLast) return false;
    int ch_int = word[0] - '0';
    return this->child[ch_int]->validate(word+1);
}

int main () {
    // set fstream
    fin.open("resource/input.txt");
    fout.open("resource/output.txt");    
 
    int T;
    cin >> T;

    for ( int t = 0 ; t < T; t++ ) {
        // input 
        int N;
        int i;

        cin >> N;
        
        char a[N][11];

        for ( i = 0; i < N; i++ ) {
            cin >> a[i];
        } 

        // make trie
        Node* rootNode = new Node();
        
        for ( i = 0; i < N; i++ ) {
            rootNode->insert(a[i]);
        }

        for ( i = 0 ; i < N ; i++ ) {
            if ( !rootNode->validate(a[i]) ) {
                cout << "NO" << endl;
                break;
            }
        }

        if ( i == N ) cout << "YES" << endl;
    }

    return 0;
}