Add GetCount() and IsEmpty()

Lexdysic · Mar 16, 2018 · c5c193f · c5c193f
1 parent d141c8a
commit c5c193f
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 19 deletions.
diff --git a/include/key_aware/key_aware.h b/include/key_aware/key_aware.h
@@ -49,13 +49,17 @@ class Trie {
 
     void Search (const StringView & prefix, std::set<std::string> * output);
 
+    inline int32_t GetCount () const { return m_count; }
+    inline bool IsEmpty () const { return !m_count; }
+
 private:
     struct Node {
         std::string            value;
         std::map<char, Node *> children;
     };
 
-    Node m_root;
+    Node    m_root;
+    int32_t m_count = 0;
 };
 
 

diff --git a/src/trie.cpp b/src/trie.cpp
@@ -24,36 +24,34 @@ void Trie::Add (const StringView & str) {
     }
 
     node->value = std::string(str.Ptr(), (size_t)str.Length());
+    m_count++;
 }
 
 void Trie::Search (const StringView & prefix, std::set<std::string> * output) {
     Node * node = &m_root;
 
-    int32_t i = 0;
-    for (; i < prefix.Length(); ++i) {
+    for (int32_t i = 0; i < prefix.Length(); ++i) {
         auto it = node->children.find(prefix[i]);
-        if (it != node->children.end()) {
-            node = it->second;
-        }
-        else {
-            break;
+        if (it == node->children.end()) {
+            return;
         }
+
+        node = it->second;
     }
 
-    if (i == prefix.Length()){
-        std::stack<Node *> stack;
-        stack.push(node);
+    std::stack<Node *> stack;
+    stack.push(node);
 
-        while (!stack.empty()) {
-            node = stack.top();
-            stack.pop();
+    while (!stack.empty()) {
+        node = stack.top();
+        stack.pop();
 
-            for (const auto & kv : node->children) {
-                stack.push(kv.second);
-            }
+        for (const auto & kv : node->children) {
+            stack.push(kv.second);
+        }
 
-            if (node->value.length())
-                output->emplace(node->value);
+        if (node->value.length()) {
+            output->emplace(node->value);
         }
     }
 }