diff --git a/lab2/dictionary.cc b/lab2/dictionary.cc
index 7bfd397..ac522ff 100644
--- a/lab2/dictionary.cc
+++ b/lab2/dictionary.cc
@@ -1,4 +1,5 @@
 #include "dictionary.h"
+#include "edit_distance.h"
 #include "word.h"
 #include <algorithm>
 #include <fstream>
@@ -56,19 +57,12 @@ void Dictionary::add_trigram_suggestions(std::vector<std::string> &suggestions,
 
 void Dictionary::rank_suggestions(std::vector<std::string> &suggestions,
                                   const std::string &word) const {
-    // Get trigrams of the input word
-    Word input_word(word);
-    const std::vector<std::string> &input_trigrams = input_word.get_triagrams();
-
-    // Sort suggestions based on the number of matching trigrams
+    // Sort suggestions based on the levenshtein distance
     std::sort(suggestions.begin(), suggestions.end(),
               [&](const std::string &a, const std::string &b) {
-                  Word word_a(a);
-                  Word word_b(b);
-                  unsigned int match_a = word_a.get_matches(input_trigrams);
-                  unsigned int match_b = word_b.get_matches(input_trigrams);
-                  return match_a >
-                         match_b; // Sort in descending order of match count
+                  auto dist_a = edit_distance(a, word);
+                  auto dist_b = edit_distance(b, word);
+                  return dist_a < dist_b;
               });
 }