hihocoder 1036 (Trie图)

题目链接

AC代码：

#include<set>
#include<map>
#include<cmath>
#include<queue>
#include<cstdio>
#include<vector>
#include<string>
#include<cstdlib>
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;

#define mem(a, b) (memset(a, b, sizeof(a)))
#define pb push_back
#define all(v) v.begin(), v.end()
#define rall(v) v.rbegin(), v.rend()
#define rep(i, m) for (int i = 0; i < (int)(m); i++)
#define rep2(i, n, m) for (int i = n; i < (int)(m); i++)
typedef long long LL;
typedef pair<int, int> pii;

const int oo = (int) 1e9;
const double PI = 2 * acos(0);
const double eps = 1e-9;
const int MAX_N = 1000010;

struct TrieNode {
    TrieNode *next[26]; // 若有子节点，则为子节点，否则为该节点后缀指向节点
    TrieNode *suffix;
    bool marked;
    TrieNode() {
        for (int i = 0; i < 26; ++i) {
            next[i] = NULL;
        }
        suffix = NULL;
        marked = false;
    }
};

struct TrieGragh {
    TrieNode *root;
    TrieGragh() { // 略去析构函数
        root = new TrieNode;
        root->suffix = root;
    }

    void insert(const char *s) {
        int pos = 0;
        TrieNode *nodePos = root;
        int sz = strlen(s);
        while (pos < sz) {
            int id = s[pos] - 'a';
            if (nodePos->next[id] == NULL) {
                nodePos->next[id] = new TrieNode;
            }
            nodePos = nodePos->next[id];
            ++pos;
        }
        nodePos->marked = true;
    }

    void generateNext() {
        queue<TrieNode *> nodeQueue;
        for (int i = 0; i < 26; ++i) {
            if (root->next[i] != NULL) { // 初始化
                nodeQueue.push(root->next[i]); // 子节点
                root->next[i]->suffix = root;
                if (root->marked)
                    root->next[i]->marked = true;
            } else {
                root->next[i] = root;
            }
        }

        while (!nodeQueue.empty()) {
            TrieNode *node = nodeQueue.front();
            nodeQueue.pop();

            // 对于该节点，每个子节点压入队列并求后缀
            // 若父节点通过ch指向子节点，则子节点后缀 = 父节点后缀通过ch指向的节点
            for (int i = 0; i < 26; ++i) {
                TrieNode *child = node->next[i];
                if (child == NULL) {
                    node->next[i] = node->suffix->next[i];
                    continue;
                }
                child->suffix = node->suffix->next[i];
                if (child->suffix->marked) // 后缀节点为已标记节点的节点也要标记
                    child->marked = true;
                if (!child->marked) // 已标记节点的子节点无需再看
                    nodeQueue.push(child);
            }
        }
    }
};

bool match(TrieGragh *trieGragh, char *article) {
    int pos = 0;
    TrieNode *nodePos = trieGragh->root;
    if (nodePos->marked)
        return true;
    int sz = strlen(article);
    while (pos < sz) {
        int id = article[pos] - 'a';
        if (id < 0 || id > 26) {
            nodePos = trieGragh->root;
            ++pos;
            continue;
        }
        nodePos = nodePos->next[id];
        if (nodePos->marked)
            return true;
        ++pos;
    }
    return false;
}

char str[MAX_N];
int main(void) {
    int N;
    scanf("%d", &N);
    TrieGragh *trieGragh = new TrieGragh;
    for (int i = 0; i < N; ++i) {
        scanf("%s", str);
        trieGragh->insert(str);
    }

    trieGragh->generateNext();
    scanf("%s", str);
    if (match(trieGragh, str))
        printf("YES
");
    else
        printf("NO
");

    return 0;
}