Wednesday 2 November 2011

Sample TRIE implementation in Java

Tried my hands at implementing a TRIE data structure in Java. There are 2 classes posted here, Trie itself and TrieNode which represents a node in a Trie.

Trie class
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;

/**
 * TRIE data structure supporting basic dictionary operations.
 */
public class Trie {

    private final TrieNode root;

    /**
     * Creates a new empty TRIE object.
     */
    public Trie() {
        root = new TrieNode();
    }


    /**
     * Inserts the specified key into this Trie object.
     * @param key
     */
    public void insert(String key) {
        TrieNode currNode = root;
        for (char c : key.toCharArray()) {
            TrieNode child = currNode.traverse(c);
            if (child == null) {
                currNode = currNode.addEdge(c);
            } else {
                currNode = child;
            }
        }
        currNode.setKey(key);
    }

    /**
     * Returns all the keys in the Trie which start with the specified prefix.
     * @param prefix
     * @return
     */
    public List<String> search(String prefix) {
        TrieNode currNode = root;
        for (char c : prefix.toCharArray()) {
            TrieNode child = currNode.traverse(c);
            if (child == null) {
                return Collections.emptyList();
            } else {
                currNode = child;
            }
        }
        List<String> matches = new ArrayList<String>();
        preorderTraverse(currNode, matches);
        return matches;
    }

    /**
     * Does preorder traversal of Trie and add retrieved keys in the specified results list.
     * @param currNode
     * @param results
     */
    private void preorderTraverse(TrieNode currNode, List<String> results) {
        if (currNode == null) return;
        if (currNode.getKey() != null) {
            results.add(currNode.getKey());
        }
        Iterator<TrieNode> children = currNode.getChildren();
        if (children != null) {
            while (children.hasNext()) {
                preorderTraverse(children.next(), results);
            }
        }
    }


    public static void main(String[] args) {
        Trie t = new Trie();
        t.insert("vino");
        t.insert("vinod");
        t.insert("vin");
        t.insert("jyo");
        t.insert("jyotsna");
        t.insert("jyot");
        t.insert("jyots");
        t.insert("jyotsn");
        t.insert("joe");

        System.out.println(t.search("vino"));
        System.out.println(t.search("j"));
        System.out.println(t.search("jy"));
        System.out.println(t.search("joe"));

        System.out.println(t.search("bhalblah"));
    }
}



TrieNode class
import java.util.*;

/**
 * A Trie Node
 */
class TrieNode {


    /**
     * The key stored in this node if any.
     */
    private String key;

    /**
     * the outgoing edges of this node, implemented as a sorted map of character to the child node.
     */
    private SortedMap<Character, TrieNode> edges;


    TrieNode addEdge(char c) {
        if (edges == null) {
            edges = new TreeMap<Character, TrieNode>();
        }
        TrieNode childNode = new TrieNode();
        edges.put(c, childNode);
        return childNode;
    }

    TrieNode traverse(char c) {
        return (edges == null) ? null : edges.get(c);
    }

    TrieNode deleteEdge(char c) {
        return (edges == null) ? null : edges.remove(c);
    }

    Iterator<TrieNode>  getChildren() {
        return (edges == null) ? null : edges.values().iterator();
    }

    void setKey(String key) {
        this.key = key;

    }

    String getKey() {
        return key;
    }

    public int getChildrenCnt() {
        return edges == null ? 0 : edges.size();
    }
}


1 comment:

  1. cool stuff....
    any optimization can you think of while doing the search.

    ReplyDelete