leetcode 212: Word Search II

Use the data structure Trie to do pruning while DFS. But you need to be careful about the way you implement the Trie. I tried the map<char, TrieNode*> to save all connections and got a TLE. So I then used TrieNode* next[26] and got accepted.

class TrieNode {
public:
    // Initialize your data structure here.
    TrieNode() {
        for(int i=0;i<26;i++)
            next[i]=NULL;
    }
    TrieNode(char _c) {
        c=_c;
        isword=0;
        for(int i=0;i<26;i++)
            next[i]=NULL;
    }
    char c;
    bool isword;
    TrieNode* next[26];
};

class Trie {
public:
    Trie() {
        root = new TrieNode();
    }

    // Inserts a word into the trie.
    void insert(string word) {
        int len=word.length();
        TrieNode *curr=root;
        for(int i=0;i<len;i++)
        {
            int idx=word[i]-'a';
            if(curr->next[idx]==NULL)
                curr->next[idx]=new TrieNode(word[i]);
            curr=curr->next[idx];
        }
        curr->isword=1;
    }

    // Returns if the word is in the trie.
    bool search(string word) {
        int len=word.length();
        TrieNode *curr=root;
        for(int i=0;i<len;i++)
        {
            int idx=word[i]-'a';
            if(curr->next[idx]==NULL)
                return false;
            curr=curr->next[idx];
        }
        return curr->isword;
    }

    // Returns if there is any word in the trie
    // that starts with the given prefix.
    bool startsWith(string prefix) {
        int len=prefix.length();
        TrieNode *curr=root;
        for(int i=0;i<len;i++)
        {
            int idx=prefix[i]-'a';
            if(curr->next[idx]==NULL)
                return false;
            curr=curr->next[idx];
        }
        return true;
    }

private:
    TrieNode* root;
};
class Solution {
public:
    vector<string> findWords(vector<vector<char>>& board, vector<string>& words) {
        Trie tree;
        int m=board.size();
        int n=board[0].size();
        int wordnum=words.size();
        for(int i=0;i<wordnum;i++)
            tree.insert(words[i]);
        set<string> st;
        string s;
        for(int i=0;i<m;i++)
            for(int j=0;j<n;j++)
                dfs_search(i,j,board,tree,st,s);
        vector<string> res;
        set<string>::iterator it;
        for(it=st.begin();it!=st.end();it++)
            res.push_back(*it);
        return res;
    }
    void dfs_search(int i,int j,vector<vector<char> >& board,Trie& tree,set<string>& st,string& curr)
    {
        if(tree.startsWith(curr+board[i][j]))
        {
            curr.push_back(board[i][j]);
            if(tree.search(curr))
                st.insert(curr);
            char temp=board[i][j];
            board[i][j]='*';
            if(i!=0&&board[i-1][j]!='*')
                dfs_search(i-1,j,board,tree,st,curr);
            if(j!=0&&board[i][j-1]!='*')
                dfs_search(i,j-1,board,tree,st,curr);
            if(i!=board.size()-1&&board[i+1][j]!='*')
                dfs_search(i+1,j,board,tree,st,curr);
            if(j!=board[0].size()-1&&board[i][j+1]!='*')
                dfs_search(i,j+1,board,tree,st,curr);
            board[i][j]=temp;
            curr.pop_back();
        }
    }
};