백준 9250 문자열 집합 판별 C++

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

9250번: 문자열 집합 판별

아호 코라식 개념에 대해 공부해볼 수 있는 문제였습니다.

 

아호 코라식이란, KMP 확장 버전으로, 길이 N인 문자열과 길이 M1, M2, .., 인 문자열 간의 매칭을 O(N+M1+M2+...) 로 진행합니다. 즉, KMP는 1대1 매칭이라면 아호 코라식은 1대다 매칭 알고리즘입니다.

 

(참고 : https://ansohxxn.github.io/algorithm/ahocorasick/) 

 

먼저 주어진 단어들에 대해 Trie를 이용하여 Fail 함수를 적용합니다. 그리고 이후 제시되는 문자열에 대해 KMP를 진행하여 단어의 포함 여부를 판단합니다.

 

코드는 다음과 같습니다.

 

#include <iostream>
#include <string>
#include <map>
#include <queue>

using namespace std;

struct Trie {
    bool isEnd; // 노드 끝인지 여부
    map<char, Trie*> child; // 자식 노드 
    Trie* fail; // 실패 노드 
    
    Trie() {
        isEnd = false;
        fail = nullptr;
    }
    
    void Insert(string s) {
        Trie* now = this;
        for(int i = 0; i < s.size(); i++) {
            // 더 이상 만족하는 char 없음
            if(now->child.find(s[i]) == now->child.end())
                now->child[s[i]] = new Trie;
            now = now->child[s[i]];
            
            // 끝까지 도달
            if(i == s.size() - 1) {
                now->isEnd = true;
            }
        }
    }
    
    // 실패 노드 전처리 
    void Fail() {
        Trie* root = this;
        queue<Trie*> q;
        q.push(root);
        
        // bfs로 탐색하면서 실패 노드 링크 
        while(!q.empty()) {
            Trie* now = q.front();
            q.pop();
            
            for(auto& ch : now->child) {
                Trie* next = ch.second;
                
                // 루트의 실패 노드는 루트 
                if(now == root)
                    next->fail = root;
                    
                else {
                    // 현재 탐색 노드의 실패 노드
                    Trie* prev = now->fail;
                    
                    // 실패 노드를 기준으로 일치하는 char 나올 때까지 거슬러 올라감 
                    while(prev != root && (prev->child.find(ch.first) == prev->child.end()))
                        prev = prev->fail;
                    
                    // 실패 노드를 결정한 실패 노드의 자식 노드로 연결 
                    if(prev->child.find(ch.first) != prev->child.end())
                        prev = prev->child[ch.first];
                    
                    next->fail = prev;
                }
                
                // 실패 노드가 끝인 경우, next 노드가 해당 실패 노드로 갈 수 있으므로 isEnd=true
                if(next->fail->isEnd)
                    next->isEnd = true;
                    
                q.push(next);
            }
        }
    }
};

bool KMP(string s, Trie* root) {
    // s : 검색할 string 
    // root : 검색될 단어들로 미리 만든 Trie 
    
    // res : s에서 Trie에 있는 단어 중 찾았는 지 유무 
    bool res = false; 
    Trie* now = root;
    for(int i = 0; i < s.size(); i++) {
        char c = s[i];
        
        // c에 대해서 fail 위치까지 이동 
        while(now != root && (now->child.find(c) == now->child.end()))
            now = now->fail;
        
        // 다음 노드로 이동 가능 
        if(now->child.find(c) != now->child.end())
            now = now->child[c];
        
        // 해당 단어 찾은 경우 
        if(now->isEnd) {
            res = true;
            break;
        }
    }
    
    return res;
}

int main()
{
    ios_base :: sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    
    int N, Q;
    string s;
    bool res;
    
    cin >> N;
    
    Trie* root = new Trie;
    for(int i = 0; i < N; i++) {
        cin >> s;
        root->Insert(s);
        s.clear();
    }
    
    root->Fail();
    
    cin >> Q;
    for(int i = 0; i < Q; i++) {
        cin >> s;
        res = KMP(s, root);
        s.clear();
        
        if(res) 
            cout << "YES\n";
        else
            cout << "NO\n";
    }
}