【PTA数据结构 | C语言版】构建后缀树

本专栏持续输出数据结构题目集，欢迎订阅。

题目

请编写程序，利用后缀树判断任一字符串是否给定字符串 s 的后缀。
当然不用后缀树也可以解决，不过本题旨在训练后缀树的实现，所以建议读者尝试用后缀树解决这个问题。

输入格式：
输入首先给出一个仅由小写英文字母组成、长度不超过 1000 的字符串，以回车结尾。以上为给定字符串 s。
接下来给出查询请求。首先给出一个正整数 n（≤1000），随后 n 行，每行给出一个仅由小写英文字母组成、长度不超过 1000 的字符串，以回车结尾。

输出格式：
对每个查询的字符串，如果它是给定字符串 s 的后缀，则在一行中输出 yes；否则输出 no。

输入样例：
binarytree
5
binary
tree
binarytree
trie
narytree

输出样例：
no
yes
yes
no
yes

代码

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define MAX_LENGTH 1001
#define ALPHABET_SIZE 26

// 后缀树节点结构
typedef struct SuffixTreeNode {
    struct SuffixTreeNode *children[ALPHABET_SIZE];
    int isEndOfSuffix; // 标记是否是某个后缀的结束
} SuffixTreeNode;

// 创建新节点
SuffixTreeNode* newNode() {
    SuffixTreeNode* node = (SuffixTreeNode*)malloc(sizeof(SuffixTreeNode));
    if (node == NULL) exit(EXIT_FAILURE);
    for (int i = 0; i < ALPHABET_SIZE; i++)
        node->children[i] = NULL;
    node->isEndOfSuffix = 0;
    return node;
}

// 插入后缀到后缀树
void insert(SuffixTreeNode* root, const char *suffix) {
    SuffixTreeNode* current = root;
    for (int i = 0; suffix[i] != '\0'; i++) {
        int index = suffix[i] - 'a';
        if (!current->children[index])
            current->children[index] = newNode();
        current = current->children[index];
    }
    current->isEndOfSuffix = 1; // 标记为后缀的结束
}

// 构建后缀树
SuffixTreeNode* buildSuffixTree(const char *s) {
    SuffixTreeNode* root = newNode();
    int len = strlen(s);
    // 插入所有后缀
    for (int i = 0; i < len; i++)
        insert(root, s + i);
    return root;
}

// 检查字符串是否是后缀
int isSuffix(SuffixTreeNode* root, const char *str) {
    SuffixTreeNode* current = root;
    for (int i = 0; str[i] != '\0'; i++) {
        int index = str[i] - 'a';
        if (!current->children[index])
            return 0;
        current = current->children[index];
    }
    return current->isEndOfSuffix; // 判断是否是后缀的结束
}

// 释放后缀树内存
void freeSuffixTree(SuffixTreeNode* node) {
    if (node == NULL) return;
    for (int i = 0; i < ALPHABET_SIZE; i++)
        if (node->children[i])
            freeSuffixTree(node->children[i]);
    free(node);
}

int main() {
    char s[MAX_LENGTH];
    if (fgets(s, MAX_LENGTH, stdin) == NULL) return 1;
    // 移除换行符
    size_t len = strlen(s);
    if (len > 0 && s[len - 1] == '\n')
        s[len - 1] = '\0';

    // 构建后缀树
    SuffixTreeNode* root = buildSuffixTree(s);

    int n;
    scanf("%d", &n);
    getchar(); // 消耗掉scanf后的换行符

    // 处理查询
    for (int i = 0; i < n; i++) {
        char query[MAX_LENGTH];
        if (fgets(query, MAX_LENGTH, stdin) == NULL) continue;
        // 移除换行符
        len = strlen(query);
        if (len > 0 && query[len - 1] == '\n')
            query[len - 1] = '\0';

        // 检查是否是后缀并输出结果
        if (isSuffix(root, query))
            printf("yes\n");
        else
            printf("no\n");
    }

    return 0;
}