ProgMe

Untitled

Dec 30th, 2023
872
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 2.77 KB | None | 0 0
  1. #include <bits/stdc++.h>
  2.  
  3. using namespace std;
  4.  
  5. class Trie {
  6.     struct Node {
  7.         int cnt = 0;
  8.         int pref = 0;
  9.  
  10.         Node* pr = nullptr;
  11.  
  12.         map<char, Node *> next;
  13.     };
  14.  
  15.     Node* root = nullptr;
  16.  
  17.     void print(Node* cur, string s);
  18.  
  19.     string kth(int idx, Node* cur, string s);
  20. public:
  21.     Trie();
  22.  
  23.     void insert(const string& s);
  24.  
  25.     void erase(const string& s);
  26.  
  27.     bool check(const string& s);
  28.  
  29.     void print();
  30.  
  31.     string kth(int idx);
  32. };
  33.  
  34. Trie::Trie() {
  35.     root = new Node();
  36. }
  37.  
  38. void Trie::insert(const std::string& s) {
  39.     Node* cur = root;
  40.  
  41.     for (auto i : s) {
  42.         ++cur->pref;
  43.         if (!cur->next[i]) {
  44.             cur->next[i] = new Node();
  45.             cur->next[i]->pr = cur;
  46.         }
  47.         cur = cur->next[i];
  48.     }
  49.  
  50.     ++cur->cnt;
  51.     ++cur->pref;
  52. }
  53.  
  54. void Trie::erase(const string& s) {
  55.     Node* cur = root;
  56.  
  57.     for (auto i : s) {
  58.         if (!cur->next[i]) {
  59.             return;
  60.         }
  61.         cur = cur->next[i];
  62.     }
  63.  
  64.     if (cur->cnt) {
  65.         --cur->cnt;
  66.         --cur->pref;
  67.         while (cur->pr) {
  68.             cur = cur->pr;
  69.             --cur->pref;
  70.         }
  71.     }
  72. }
  73.  
  74. bool Trie::check(const string& s) {
  75.     Node* cur = root;
  76.  
  77.     for (auto i : s) {
  78.         if (!cur->next[i]) {
  79.             return false;
  80.         }
  81.         cur = cur->next[i];
  82.     }
  83.  
  84.     return cur->cnt;
  85. }
  86.  
  87. void Trie::print() {
  88.     print(root, "");
  89. }
  90.  
  91. void Trie::print(Node* cur, string s) {
  92.  
  93.     for (auto [c, to] : cur->next) {
  94.         if (to) {
  95.             s.push_back(c);
  96.             print(to, s);
  97.             s.pop_back();
  98.         }
  99.     }
  100.  
  101.     if (cur->cnt) {
  102.         cout << s << '\n';
  103.     }
  104. }
  105.  
  106. string Trie::kth(int idx) {
  107.     return kth(idx, root, "");
  108. }
  109.  
  110. string Trie::kth(int idx, Node* cur, string s) {
  111.     if (cur->cnt > idx) {
  112.         return s;
  113.     }
  114.     idx -= cur->cnt;
  115.  
  116.     for (auto [c, to] : cur->next) {
  117.         if (to->pref > idx) {
  118.             s.push_back(c);
  119.             return kth(idx, to, s);
  120.         }
  121.         idx -= to->pref;
  122.     }
  123.  
  124.     return "-1";
  125. }
  126.  
  127. int main() {
  128.     Trie t;
  129.  
  130.     int q;
  131.     cin >> q;
  132.     while (q--) {
  133.         int tp;
  134.         cin >> tp;
  135.  
  136.         if (tp == 1) {
  137.             string s;
  138.             cin >> s;
  139.            t.insert(s);
  140.         } else if (tp == 2) {
  141.             string s;
  142.             cin >> s;
  143.             t.erase(s);
  144.         } else if (tp == 3) {
  145.             string s;
  146.             cin >> s;
  147.             cout << t.check(s) << '\n';
  148.             cout << '\n';
  149.         } else if (tp == 4) {
  150.             t.print();
  151.             cout << '\n';
  152.         } else {
  153.             int idx;
  154.             cin >> idx;
  155.             cout << t.kth(idx) << '\n';
  156.             cout << '\n';
  157.         }
  158.     }
  159. }
Advertisement
Add Comment
Please, Sign In to add comment