【算法学习笔记】57. 前缀树 字典序优化技巧 STL学习 SJTU OJ 1366 前缀匹配

Description

给出一个总字符个数大小不超过1，000，000的字典（这个字典的单词顺序不为字典序）和不超过1000个长度不超过1000的前缀，输出字典中匹配该前缀，字典序为K_i的单词在字典中的位置。

所有单词都为小写字母。

Input Format

第一行： 两个整数N，M，分别表示字典中的单词个数和需要查询的前缀数。

接下来N行；每行一个字符串，表示字典中的单词。

接下来M行，每行一个K_i， P_i, P_i表示查询的前缀,K_i表示满足这个前缀在字典序中的位置。

Output Format

M行整数，表示符合条件的单词在字典中的位置。如果没有单词满足要求，输出-1

样例解释，符合前缀'a'的单词有{'az'， 'ay'}，按字典序排序为{'ay', 'az'},字典序第2的是az，在字典的第一个位置。

Input Sample

3 2
az
ay
b
2 a
3 a

Output Sample

1
-1

最简单的思路就是按照Trie树的结构来进行建树，然后查询时DFS（每一行逆序扫描）来查找以该前缀为前缀的单词，找到第k个然后返回原始id即可
但是这样的查询当某一个前缀的单词比较多而且k比较大的时候 建树和查询的时间都会很长
所以 完全可以不用Trie树，直接用字典序排序 + 二分查找即可。（标程的思想）
想在Trie上优化的话，也是可以的，先排成字典序，然后Trie查找的时候找到第一个单词，返回字典序id 判断 id+k-1的字典序单词是否还有该前缀。

我的代码如下：

#include <iostream>
#include <cstring>
#include <algorithm>
#include <cstdio>
using namespace std;
const int MaxN = 26;

//单词的结构
struct Word
{
    char word[100];//单词的内容
    int word_len;//单词的长度
    int ori_id;//单词的原始ID
};

Word dictionary[1000000];//存储字典 后来会排成字典序


//树的节点
struct TrieNode
{
    bool isEnd; //是否是一个单词的结尾字符
    int sonCount; //子节点的个数
    TrieNode* son[MaxN];//子节点指针列表 最多只有26个
    int id;//以此节点字符为结尾的单词 在排序后的字典中的位置
    TrieNode(){
        isEnd = false;
        sonCount = 0;
        id = -1;
        for (int i = 0; i < MaxN; ++i)
            son[i] = NULL;
    }
};

TrieNode* stack[1000000] = {0};//用于DFS的堆栈

TrieNode* root;//根节点 代表整个字典树

//插入一个词 给出词的内容 词的长度 以及 词在单词列表中的位置
void insert(char* word, int len, int id){
    TrieNode* curNode = root;//以根节点为起点
    //cout<<root->sonCount<<endl;
    for (int i = 0; i < len; ++i)
    {
        //判断word[i]是否在curNode的子节点列表里
        int index = word[i]-'a'; //index 0~25
        
        if(curNode->son[index]==NULL){//没有此字符
            curNode->son[index] = new TrieNode();
            //个数加一
            curNode->sonCount++;
        }
        
        if(i==len-1){//单词的结尾
            curNode->son[index]->isEnd = true;
            curNode->son[index]->id = id;
        }
        
        curNode = curNode->son[index];
    }
}

bool cmp_int (const int& a, const int& b){
    return a < b;
}
//字典序排序
bool cmp_word(const Word& a, const Word& b){
    int todo = a.word_len < b.word_len ? a.word_len : b.word_len;
    for (int i = 0; i < todo; ++i)
    {
        if(a.word[i]==b.word[i])
            continue;
        return a.word[i] < b.word[i];
    }
    //此时还未return说明两个单词互相包含 把少的那个放在前面
    return a.word_len < b.word_len;
}

//根据前缀查询 第cnt个单词 返回在字典中的id
int find(int cnt, char* prefix , int len){
    TrieNode* curNode = root;
    
    for (int i = 0; i < len; ++i)
    {
        int cid = prefix[i]-'a';
        if(curNode->son[cid] != NULL){
            curNode = curNode->son[cid];
        }
    }
    //现在已经找到了前缀的最后一个字符的节点
    int top = 0;
    stack[top++] = curNode;
    //DFS去找第一个该前缀的词 
    while(top > 0){
        TrieNode* toCheck =  stack[--top];
        if(toCheck->isEnd){
            //立刻返回
            return toCheck->id;
        }else{
            int haveAdded = 0;
            for (int i = MaxN-1; i >=0 ; --i)
            {
                if(toCheck->son[i] != NULL){
                    haveAdded++;
                    stack[top++] = toCheck->son[i];
                    if(haveAdded==toCheck->sonCount)
                        break;
                }
            }
        }
    }
    return -1;
}

//判断单词word是否是以prefix为前缀的
bool isIn(Word& word, char* prefix , int len){
    if(word.word_len < len)
        return false;
    for(int i=0; i < len; i++){
        if(word.word[i] != prefix[i])
            return false;
    }
    return true;
}



int main(int argc, char const *argv[])
{
    int N,M;
    cin>>N>>M; 
    root = new TrieNode();//初始化根节点
    
    for (int i = 1; i <= N; ++i)
    {
        scanf("%s",dictionary[i].word); 
        dictionary[i].word_len = strlen(dictionary[i].word);
        dictionary[i].ori_id = i;
        
    }
    
    sort(dictionary+1, dictionary+N+1, cmp_word);
    
    for (int i = 1; i <= N; ++i)
        insert(dictionary[i].word, dictionary[i].word_len, i);

    char pre[1000+10];
    for (int i = 0; i < M; ++i)
    {
        int k = 0;
        scanf("%d %s",&k,pre);
        //找到以该前缀为前缀的第一个单词的排序后ID + k-1 判断是否是
        int first = find(k,pre,strlen(pre));

        if(first+k-1 <=N and isIn(dictionary[first+k-1],pre,strlen(pre))){
            cout<<dictionary[first+k-1].ori_id<<endl;
        }else
            cout<<-1<<endl;
    }
    return 0;
}

/*
4 2
ad
as
ase
asfg
2 as
1 ad

 */

标程如下：

#include <iostream>
#include <string>
#include <vector>
#include <algorithm>
using namespace std;
 
int n, m;
vector< pair<string, int> > dict;
 
//判断word是否含有pref前缀 果然stl好快
bool match(string& pref, string& word) {
    if (pref.size() > word.size()) return false;
    return word.substr(0, pref.size()) == pref;
}
 
int main() {
    cin >> n >> m;
    for (int i = 0; i < n; i++) {
        string s;
        cin >> s;
        dict.push_back(make_pair(s, i));//pair的单词和id
    }
    //排字典序 因为string有重载的大于号 所以不用加入函数
    sort(dict.begin(), dict.end());

    for (int i = 0; i < m; i++) {
        int k;
        string pref;
        cin >> k >> pref;
        k--;//因为id从0开始
        //stl的二分查找 
        int pos = lower_bound(dict.begin(), dict.end(), make_pair(pref, 0)) - dict.begin();
        // Ignore whether pos even matches at all
        int poss = pos + k;
        if (poss < dict.size() && match(pref, dict[poss].first)) {
            cout << dict[poss].second + 1 << '
';
        } else {
            cout << -1 << '
';
        }
    }
}