Algorithms : complete 04-algorithms_on_strings 01-suffix_trees

1 files changed, 82 insertions, 13 deletions

diff --git a/04-algorithms_on_strings/01-suffix_trees/02-trie_matching/trie_matching.cpp b/04-algorithms_on_strings/01-suffix_trees/02-trie_matching/trie_matching.cpp
index a59a0db..2517d81 100644
--- a/04-algorithms_on_strings/01-suffix_trees/02-trie_matching/trie_matching.cpp
+++ b/04-algorithms_on_strings/01-suffix_trees/02-trie_matching/trie_matching.cpp
@@ -23,27 +23,96 @@ struct Node
 	{
 	    return (next[0] == NA && next[1] == NA && next[2] == NA && next[3] == NA);
 	}
+
+    void print()
+    {
+        if (isLeaf())
+            cout << "---- " << endl;
+        else
+            cout << "A:" << next[0] << " C:" << next[1]
+            << " G:" << next[2] << " T:" << next[3] << endl;
+    }
 };
 
-int letterToIndex (char letter)
+struct Trie
 {
-	switch (letter)
-	{
-		case 'A': return 0; break;
-		case 'C': return 1; break;
-		case 'G': return 2; break;
-		case 'T': return 3; break;
-		default: assert (false); return -1;
-	}
-}
+    vector<Node> nodes;
+
+    Trie(const vector <string>& patterns)
+    {
+        int v = 1;
+        nodes.push_back(Node());
+
+        for (int i = 0; i < patterns.size(); i++) {
+            int n = 0;
+            const string& pattern = patterns[i];
+            for (int j = 0; j < pattern.size(); j++) {
+                int c = letterToIndex(pattern[j]);
+                int next = nodes[n].next[c];
+                if (next != NA) {
+                    n = next;
+                } else {
+                    nodes[n].next[c] = v;
+                    n = v;
+                    nodes.push_back(Node());
+                    v += 1;
+                }
+            }
+        }
+    }
+
+    int letterToIndex (char letter)
+    {
+        switch (letter)
+        {
+            case 'A': return 0; break;
+            case 'C': return 1; break;
+            case 'G': return 2; break;
+            case 'T': return 3; break;
+            default: assert (false); return -1;
+        }
+    }
+
+    void print()
+    {
+        for (int i = 0; i < nodes.size(); i++) {
+            cout << i << " " ;
+            nodes[i].print();
+        }
+    }
+
+    bool prefixMatch(const string& text)
+    {
+        int n = 0;
+
+        for( int i = 0; i < text.size(); i++) {
+            int next = nodes[n].next[letterToIndex(text[i])];
+            if (next == NA)
+                break;
+            n = next;
+        }
+
+        return nodes[n].isLeaf();
+    }
+
+    void match(const string& text, vector<int>& result)
+    {
+        for (int i = 0; i < text.size(); i++) {
+            if (prefixMatch(&text[i]))
+                result.push_back(i);
+        }
+    }
+};
 
 vector <int> solve (const string& text, int n, const vector <string>& patterns)
 {
-	vector <int> result;
+    vector <int> result;
 
-	// write your code here
+    Trie t(patterns);
+    /* t.print(); */
+    t.match(text, result);
 
-	return result;
+    return result;
 }
 
 int main (void)