pylucene 3.5.0-3

[pylucene.git] / lucene-java-3.5.0 / lucene / contrib / spellchecker / src / java / org / apache / lucene / search / suggest / tst / TSTAutocomplete.java

diff --git a/lucene-java-3.5.0/lucene/contrib/spellchecker/src/java/org/apache/lucene/search/suggest/tst/TSTAutocomplete.java b/lucene-java-3.5.0/lucene/contrib/spellchecker/src/java/org/apache/lucene/search/suggest/tst/TSTAutocomplete.java
new file mode 100644 (file)
index 0000000..5235dfa
--- /dev/null
+++ b/lucene-java-3.5.0/lucene/contrib/spellchecker/src/java/org/apache/lucene/search/suggest/tst/TSTAutocomplete.java
@@ -0,0 +1,142 @@
+package org.apache.lucene.search.suggest.tst;
+
+import java.util.*;
+
+public class TSTAutocomplete {
+
+  /**
+   * Inserting keys in TST in the order middle,small,big (lexicographic measure)
+   * recursively creates a balanced tree which reduces insertion and search
+   * times significantly.
+   * 
+   * @param tokens
+   *          Sorted list of keys to be inserted in TST.
+   * @param lo
+   *          stores the lower index of current list.
+   * @param hi
+   *          stores the higher index of current list.
+   * @param root
+   *          a reference object to root of TST.
+   */
+  public void balancedTree(Object[] tokens, Object[] vals, int lo, int hi,
+          TernaryTreeNode root) {
+    if (lo > hi) return;
+    int mid = (lo + hi) / 2;
+    root = insert(root, (String) tokens[mid], vals[mid], 0);
+    balancedTree(tokens, vals, lo, mid - 1, root);
+    balancedTree(tokens, vals, mid + 1, hi, root);
+  }
+
+  /**
+   * Inserts a key in TST creating a series of Binary Search Trees at each node.
+   * The key is actually stored across the eqKid of each node in a successive
+   * manner.
+   * 
+   * @param currentNode
+   *          a reference node where the insertion will take currently.
+   * @param s
+   *          key to be inserted in TST.
+   * @param x
+   *          index of character in key to be inserted currently.
+   * @return currentNode The new reference to root node of TST
+   */
+  public TernaryTreeNode insert(TernaryTreeNode currentNode, String s,
+          Object val, int x) {
+    if (s == null || s.length() <= x) {
+      return currentNode;
+    }
+    if (currentNode == null) {
+      TernaryTreeNode newNode = new TernaryTreeNode();
+      newNode.splitchar = s.charAt(x);
+      currentNode = newNode;
+      if (x < s.length() - 1) {
+        currentNode.eqKid = insert(currentNode.eqKid, s, val, x + 1);
+      } else {
+        currentNode.token = s;
+        currentNode.val = val;
+        return currentNode;
+      }
+    } else if (currentNode.splitchar > s.charAt(x)) {
+      currentNode.loKid = insert(currentNode.loKid, s, val, x);
+    } else if (currentNode.splitchar == s.charAt(x)) {
+      if (x < s.length() - 1) {
+        currentNode.eqKid = insert(currentNode.eqKid, s, val, x + 1);
+      } else {
+        currentNode.token = s;
+        currentNode.val = val;
+        return currentNode;
+      }
+    } else {
+      currentNode.hiKid = insert(currentNode.hiKid, s, val, x);
+    }
+    return currentNode;
+  }
+
+  /**
+   * Auto-completes a given prefix query using Depth-First Search with the end
+   * of prefix as source node each time finding a new leaf to get a complete key
+   * to be added in the suggest list.
+   * 
+   * @param root
+   *          a reference to root node of TST.
+   * @param s
+   *          prefix query to be auto-completed.
+   * @param x
+   *          index of current character to be searched while traversing through
+   *          the prefix in TST.
+   * @return suggest list of auto-completed keys for the given prefix query.
+   */
+  public ArrayList<TernaryTreeNode> prefixCompletion(TernaryTreeNode root,
+          String s, int x) {
+
+    TernaryTreeNode p = root;
+    ArrayList<TernaryTreeNode> suggest = new ArrayList<TernaryTreeNode>();
+
+    while (p != null) {
+      if (s.charAt(x) < p.splitchar) {
+        p = p.loKid;
+      } else if (s.charAt(x) == p.splitchar) {
+        if (x == s.length() - 1) {
+          break;
+        } else {
+          x++;
+        }
+        p = p.eqKid;
+      } else {
+        p = p.hiKid;
+      }
+    }
+
+    if (p == null) return suggest;
+    if (p.eqKid == null && p.token == null) return suggest;
+    if (p.eqKid == null && p.token != null) {
+      suggest.add(p);
+      return suggest;
+    }
+
+    if (p.token != null) {
+      suggest.add(p);
+    }
+    p = p.eqKid;
+
+    Stack<TernaryTreeNode> st = new Stack<TernaryTreeNode>();
+    st.push(p);
+    while (!st.empty()) {
+      TernaryTreeNode top = st.peek();
+      st.pop();
+      if (top.token != null) {
+        suggest.add(top);
+      }
+      if (top.eqKid != null) {
+        st.push(top.eqKid);
+      }
+      if (top.loKid != null) {
+        st.push(top.loKid);
+      }
+      if (top.hiKid != null) {
+        st.push(top.hiKid);
+      }
+    }
+    return suggest;
+  }
+}