题解：P7475 「C.E.L.U-02」简易输入法

题意

给定词典 \(\text{U}\)，每次询问读入一个字符串 \(s\)，以及一个整数 \(x\) 对于这个字符串有以下几种情形：

设\(s_i \in \text{U}\) 且 \(s\) 为 \(s_i\) 的前缀的个数为 \(a\)。

当 \(a\ge x\) 时，请输出按照以输出次数从大到小为第一关键字，以字典序为第二关键字排序后的第 \(x\) 个 \(s_i\)，并将其输出次数加 \(1\)。

当 \(x>a>0\) 时，请输出按照以输出次数从大到小为第一关键字，以字典序为第二关键字排序后的最后一个 \(s_i\)，并将其输出次数加 \(1\)。

当 \(a=0\) 时，输出 404Error。

分析

前缀相关的统计操作，考虑使用 Trie。

带修改查询第 \(x\) 大，考虑使用平衡树。

在每个 Trie 节点上维护一棵平衡树，维护所有包含该前缀的 \(s_i\) 以及该串输出次数。

因为 \(|s_i|\leq 10\)，所以修改的时候直接模拟插入过程暴力修改每个节点上的平衡树的值。

查询的时候在尾节点上询问即可。

时间复杂度 \(O((n+q)|s|\log n)\)。

Code

因为不想重载运算符所以维护的是输出次数的相反数。

#include<bits/stdc++.h>
#include<bits/extc++.h>
using namespace std;
using namespace __gnu_pbds;

struct node
{
    node* ch[26];
    tree<pair<int, string>, null_type, less<pair<int, string>>, 
         rb_tree_tag, tree_order_statistics_node_update> tr;
    node() {memset(ch, 0, sizeof ch);}
}*rt=new node;
unordered_map<string, int> mpt;

void insert(string s)
{
    node *x=rt;
    for(auto c:s)
    {
        x->tr.insert({0, s});
        if(!x->ch[c-'a']) x->ch[c-'a']=new node;
        x=x->ch[c-'a'];
    }
    x->tr.insert({0, s});
    mpt[s]=0;
}

void modify(string s)
{
    node *x=rt;
    int v=mpt[s];
    for(auto c:s)
    {
        x->tr.erase({v, s});
        x->tr.insert({v-1, s});
        x=x->ch[c-'a'];
    }
    x->tr.erase({v, s});
    x->tr.insert({v-1, s});
    mpt[s]--;
}

string query(string s, int v)
{
    node *x=rt;
    for(auto c:s)
    {
        if(!x->ch[c-'a']) return "404Error";
        x=x->ch[c-'a'];
    }
    string p=x->tr.find_by_order(min(v, (int)x->tr.size())-1)->second;
    modify(p);
    return p;
}

int main()
{
    int n, t, x;
    string p;
    cin>>n;
    for(int i=1;i<=n;i++) cin>>p, insert(p);
    cin>>t;
    while(t--)  cin>>p>>x, cout<<query(p, x)<<'\n';
}