main.cpp

#include <iostream>
#include <fstream>
#include <set>
#include <algorithm>
#include <cctype>
using namespace std;

class Trie
{
  public:
    struct Node
    {
      Node() : final(false)
      {
        for(int i = 0;i < 26;++i)
          children[i] = 0;
      }

      ~Node()
      {
        for(int i = 0;i < 26;++i)
          if(children[i]) delete children[i];
      }

      Node*& child(char c)
      {
        return children[tolower(c)-'a'];
      }

      Node* children[26];
      bool final;
      string word;
    };

    Trie()
    {
      root = new Node;
    }

    ~Trie()
    {
      delete root;
    }

    void add(string word)
    {
      Node* curr = root;
      for(int i = 0;i < word.length();++i)
      {
        char& c = word[i];
        if(!isalpha(c)) continue;
        if(!curr->child(c))
          curr->child(c) = new Node;
        curr = curr->child(c);
      }
      curr->final = true;
      curr->word = word;
    }

    bool search(string word)
    {
      Node* curr = root;
      for(int i = 0;i < word.size();++i)
      {
        char& c = word[i];
        if(!isalpha(c)) continue;
        if(!curr->child(c)) return false;
        curr = curr->child(c);
      }
      return curr->final;
    }

    set<string> corrections(string word, int d)
    {
      set<string> res;
      corrections(word,d,root,res);
      return res;
    }
  private:
    void corrections(string word, int d, Node* curr, set<string>& res)
    {
      if(word.size() == 0)
      {
        if(curr->final)
          res.insert(curr->word);
        return;
      }

      char c = word[0];
      if(curr->child(c))
        corrections(word.substr(1),d,curr->child(c),res);

      if(d <= 0) return;

      for(int i = 0;i < 26;++i)
        if(curr->children[i])
        {
          corrections(word,d-1,curr->children[i],res); //missed a letter
          corrections(word.substr(1),d-1,curr->children[i],res); //mis-typed letter
        }

      if(word.size() < 2) return;

      //typed extra letter
      c = word[1];
      if(curr->child(c))
        corrections(word.substr(2),d-1,curr->child(c),res);

      //swapped letters
      swap(word[0],word[1]);
      corrections(word,d-1,curr,res);
    }
    Node* root;
};

int main()
{
  Trie trie;
  string tmp;
  
  cout << "indexing dictionary...";
  ifstream dict("/usr/share/dict/words");
  while(!dict.eof())
  {
    dict >> tmp;
    trie.add(tmp);
  }
  cout << " done" << endl;

  for(;;)
  {
    cout << "> ";
    cin >> tmp;
    if(tmp == "add")
    {
      cin >> tmp;
      trie.add(tmp);
      cout << "added \"" << tmp << "\"";
    }
    else if(tmp == "search")
    {
      cin >> tmp;
      cout << trie.search(tmp);
    }
    else if(tmp == "correct")
    {
      cin >> tmp;
      set<string> corrections = trie.corrections(tmp,2);
      set<string>::iterator itr;
      for(itr = corrections.begin();itr != corrections.end();++itr)
        cout << *itr << "\n";
    }
    else
      cerr << "unrecognized command";
    cout << endl;
  }
}