Friday, February 6, 2015

Trie tree -- Java

import java.util.ArrayList;


class Node {
    char c;                         // the character in the node
    boolean isEnd;                  // whether the end of the words
    int count;                      // the number of words sharing this character
    ArrayList<Node> childList;      // the child list
 
 
    public Node(char c){
        this.c = c;
        isEnd = false;
        count = 0;
        childList = new ArrayList<Node>();      
    }
 
    public Node subNode(char character){
        if(childList != null){
            for(Node eachChild : childList){
                if(eachChild.c == character){
                     return eachChild;
                }
            }
        }
        return null;
    }
}



class Trie {
   private Node root;

   public Trie(){
       root = new Node(' ');
   }
   
 
   public void insertWord(String word){
       if(searchWord(word) == true) return;
       
       Node current = root;
       for(int i = 0; i < word.length(); i++){
       
           Node child = current.subNode(word.charAt(i));
         
           if(child != null){
               current = child;           //word[i] char has existed in Trie!
           }
           else {                               //word[i] char not existed in Trie! Need to create a new Node and insert.
                int j=0;
                for(; j< current.childList.size(); ++j){
            if(current.childList.get(j).c > word.charAt(i)){
            break;
            }
            }      
                current.childList.add(j, new Node(word.charAt(i)));
           
                current = current.subNode(word.charAt(i));
           }
           current.count++;
       }
                                                       // Set isEnd to indicate end of the word, important!
       current.isEnd = true;
   }
 
 
 
   public boolean searchWord(String word){
       Node current = root;
       
       for(int i = 0; i < word.length(); i++){    
           if(current.subNode(word.charAt(i)) == null)
               return false;
           else
               current = current.subNode(word.charAt(i));
       }      
                                      //If a string exists, make sure it's a word by checking its 'isEnd' flag !       
       if (current.isEnd == true) return true;
       return false;
   }
   
 
   public void deleteWord(String word){
       if(searchWord(word) == false) return;        
     
       Node current = root;
       for(char c : word.toCharArray()) {
       
           Node child = current.subNode(c);
         
           if(child.count == 1) {
                 current.childList.remove(child);
                 return;
           }
           else {
                 child.count--;
                 current = child;
           }
       }
       current.isEnd = false;
   }

 
   public ArrayList<String> listAllWord(){                                           //DFS
   
    ArrayList<String> res = new ArrayList<String>();
    StringBuilder word = new StringBuilder();
   
    Node current = root;
    listAllWordHelp(current, word, res);
          return res;
   }

   public void listAllWordHelp(Node current, StringBuilder word, ArrayList<String> res){
    Node curr = current;
   
                   //edge case : isEnd : add that word to res.
    if(curr.isEnd){
    String subres = word.toString();
    res.add(subres);
    }
   
    for(Node eachChild : curr.childList){
   
    word.append(eachChild.c);
    curr = eachChild;
                //recursive case!
    listAllWordHelp(curr, word, res);
                //for backtrack usage!
    word.deleteCharAt(word.length()-1);
    }  
        //return case: childList.size() == 0, leaf node!!!
    return;
   }
     
 
   public static void main(String[] args) {
       Trie trie = new Trie();
     
       trie.insertWord("ball");
       trie.insertWord("balls");
       trie.insertWord("sense");
       trie.insertWord("aaa");
       trie.insertWord("aab");
       trie.insertWord("aba");
       trie.insertWord("baa");
       trie.insertWord("bab");    
       trie.insertWord("ccc");
   
     
       System.out.println(trie.searchWord("balls"));
       System.out.println(trie.searchWord("ba"));
       System.out.println(trie.searchWord("sen"));
       trie.deleteWord("balls");
       System.out.println(trie.searchWord("balls"));
       System.out.println(trie.searchWord("ball"));
     
       ArrayList<String> allWord = new ArrayList<String>();
       allWord = trie.listAllWord();
       for(String str : allWord){
        System.out.println(str);
       }    
   }
}



No comments:

Post a Comment