update light trie

hitonanode · Dec 9, 2023 · a32e95c · a32e95c
1 parent e5b7502
commit a32e95c
Showing 1 changed file with 14 additions and 4 deletions.
diff --git a/string/trie_light.hpp b/string/trie_light.hpp
@@ -1,26 +1,36 @@
 #include <string>
 #include <vector>
 
-// CUT begin
 struct Trie {
     char a_init;
     int D;
     int INVALID = -1;
     std::vector<std::vector<int>> child;
+    std::vector<int> par;
+
     using T_NODE = int;
     std::vector<T_NODE> v_info;
+
     Trie(char a_init = 'a', int D = 26)
-        : a_init(a_init), D(D), child(1, std::vector<int>(D, INVALID)), v_info(1) {}
-    void add_word(const std::string &str, T_NODE info) {
+        : a_init(a_init), D(D), child(1, std::vector<int>(D, INVALID)), par(1, -1), v_info(1) {}
+
+    int step(int now, char c) const {
+        if (now == INVALID) return INVALID;
+        return child.at(now).at(c - a_init);
+    }
+
+    int add_word(const std::string &str, T_NODE info) {
         int now = 0;
         for (auto &c : str) {
             if (child[now][c - a_init] == INVALID) {
+                par.push_back(now);
                 child[now][c - a_init] = child.size();
-                child.emplace_back(std::vector<int>(D, INVALID));
+                child.emplace_back(D, INVALID);
                 v_info.resize(child.size());
             }
             now = child[now][c - a_init];
         }
         v_info[now] += info;
+        return now;
     }
 };