Skip to content

Instantly share code, notes, and snippets.

@phyous
Created June 7, 2014 19:22
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save phyous/07acad23646db78aec86 to your computer and use it in GitHub Desktop.
Save phyous/07acad23646db78aec86 to your computer and use it in GitHub Desktop.
Having fun with Tries
package com.other.rand;
import java.util.ArrayList;
import java.util.List;
/*
Let's create a trie data structure and run some tests on it
Documentation:
- http://en.wikipedia.org/wiki/TrieTest
- http://www.geeksforgeeks.org/trie-insert-and-search/
*/
public class TrieTest {
// DATA STRUCTURE CODE
public static class Trie {
private Node mRoot;
private int mCount;
public Trie() {
mRoot = new Node('\0');
mCount = 0;
}
private class Node {
int code;
char val;
Node[] children;
private static final int ALPHABET_SIZE = 26;
Node(char val) {
this.val = val;
children = new Node[ALPHABET_SIZE];
}
}
public int getCount() {
return mCount;
}
/*
Insert a string in the trie.
Assumes s contains only lower case a-z characters
*/
public void insert(String s) {
if (s == null || s.isEmpty()) return;
Node pNode = mRoot;
for (char c : s.toCharArray()) {
int charCode = getCharInt(c);
if (pNode.children[charCode] == null) {
pNode.children[charCode] = new Node(c);
}
pNode = pNode.children[charCode];
}
if (pNode.code <= 0) pNode.code = ++mCount;
}
/*
Search for a given string in the trie
*/
public boolean search(String s) {
if (s == null || s.isEmpty()) return false;
Node pNode = mRoot;
for (char c : s.toCharArray()) {
int charCode = getCharInt(c);
if (pNode.children[charCode] == null) {
return false;
} else {
pNode = pNode.children[charCode];
}
}
return true;
}
/*
Return all strings that match prefix in the trie.
*/
public List<String> prefixSearch(String prefix) {
List<String> ret = new ArrayList<String>();
if (prefix == null) return ret;
Node pNode = mRoot;
// 1- Traverse trie to prefix end point
for (char c : prefix.toCharArray()) {
int charCode = getCharInt(c);
if (pNode.children[charCode] == null) {
return ret;
} else {
pNode = pNode.children[charCode];
}
}
// 2- Recursively build the prefix matches
StringBuilder prefixBuilder = new StringBuilder(prefix);
traverse(pNode, ret, prefixBuilder);
return ret;
}
private int getCharInt(char c) {
return ((int) c) - 97;
}
private void traverse(Node pNode, List<String> results, StringBuilder prefix) {
if (pNode.code > 0) results.add(prefix.toString());
for (Node n : pNode.children) {
if (n != null) {
prefix.append(n.val);
traverse(n, results, prefix);
prefix.deleteCharAt(prefix.length() - 1);
}
}
}
}
// TEST CODE
public static void main(String[] args) {
testInsertion();
testSearch();
testPrefixSearch();
}
private static void testPrefixSearch() {
String testCategory = "PREFIX SEARCH";
Trie trie = new Trie();
trie.insert("hip");
trie.insert("hipper");
trie.insert("hilarious");
trie.insert("cap");
List<String> matches = trie.prefixSearch("hi");
testAssert(testCategory, "Search #1 returns correct # of strings", matches.size() == 3);
boolean condition = matches.contains("hip") && matches.contains("hipper") && matches.contains("hilarious");
testAssert(testCategory, "Search #1 returns correct string values", condition);
matches = trie.prefixSearch("hip");
testAssert(testCategory, "Search #2 returns correct # of strings", matches.size() == 2);
condition = matches.contains("hip") && matches.contains("hipper");
testAssert(testCategory, "Search #2 returns correct string values", condition);
matches = trie.prefixSearch("nope");
testAssert(testCategory, "Search for non existent string returns 0", matches.size() == 0);
matches = trie.prefixSearch("");
testAssert(testCategory, "Search for empty string returns all", matches.size() == 4);
matches = trie.prefixSearch(null);
testAssert(testCategory, "Search for null string returns nothing", matches.size() == 0);
}
private static void testSearch() {
String testCategory = "SEARCH";
Trie trie = new Trie();
trie.insert("hip");
trie.insert("hipper");
trie.insert("hilarious");
trie.insert("cap");
boolean found = trie.search("hip");
testAssert(testCategory, "Search for existing string returns true", found == true);
found = trie.search("nope");
testAssert(testCategory, "Search for non-existing string returns false", found == false);
found = trie.search(null);
testAssert(testCategory, "Search for null string returns false", found == false);
found = trie.search("");
testAssert(testCategory, "Search for empty string returns false", found == false);
}
private static void testInsertion() {
String testCategory = "INSERTION";
Trie trie = new Trie();
trie.insert("");
testAssert(testCategory, "Empty strings don't get added", trie.getCount() == 0);
trie.insert(null);
testAssert(testCategory, "Null strings don't get added", trie.getCount() == 0);
trie.insert("hip");
trie.insert("hit");
trie.insert("hillarious");
trie.insert("cap");
testAssert(testCategory, "Correct strings get properly added", trie.getCount() == 4);
trie.insert("cap");
testAssert(testCategory, "Duplicated strings added don't increase the count", trie.getCount() == 4);
}
private static void testAssert(String testCategory, String testDescription, boolean condition) {
StringBuilder sb = new StringBuilder();
if (testCategory != null) sb.append(String.format("[%s] ", testCategory));
if (condition) {
sb.append("PASS - ");
} else {
sb.append("FAIL - ");
}
sb.append(testDescription);
System.out.println(sb.toString());
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment