lucene-java-3.4.0/lucene/contrib/spellchecker/src/java/org/apache/lucene/search/suggest/tst/TSTAutocomplete.java

   1 package org.apache.lucene.search.suggest.tst;
   2
   3 import java.util.*;
   4
   5 public class TSTAutocomplete {
   6
   7   /**
   8    * Inserting keys in TST in the order middle,small,big (lexicographic measure)
   9    * recursively creates a balanced tree which reduces insertion and search
  10    * times significantly.
  11    *
  12    * @param tokens
  13    *          Sorted list of keys to be inserted in TST.
  14    * @param lo
  15    *          stores the lower index of current list.
  16    * @param hi
  17    *          stores the higher index of current list.
  18    * @param root
  19    *          a reference object to root of TST.
  20    */
  21   public void balancedTree(Object[] tokens, Object[] vals, int lo, int hi,
  22           TernaryTreeNode root) {
  23     if (lo > hi) return;
  24     int mid = (lo + hi) / 2;
  25     root = insert(root, (String) tokens[mid], vals[mid], 0);
  26     balancedTree(tokens, vals, lo, mid - 1, root);
  27     balancedTree(tokens, vals, mid + 1, hi, root);
  28   }
  29
  30   /**
  31    * Inserts a key in TST creating a series of Binary Search Trees at each node.
  32    * The key is actually stored across the eqKid of each node in a successive
  33    * manner.
  34    *
  35    * @param currentNode
  36    *          a reference node where the insertion will take currently.
  37    * @param s
  38    *          key to be inserted in TST.
  39    * @param x
  40    *          index of character in key to be inserted currently.
  41    * @return currentNode The new reference to root node of TST
  42    */
  43   public TernaryTreeNode insert(TernaryTreeNode currentNode, String s,
  44           Object val, int x) {
  45     if (s == null || s.length() <= x) {
  46       return currentNode;
  47     }
  48     if (currentNode == null) {
  49       TernaryTreeNode newNode = new TernaryTreeNode();
  50       newNode.splitchar = s.charAt(x);
  51       currentNode = newNode;
  52       if (x < s.length() - 1) {
  53         currentNode.eqKid = insert(currentNode.eqKid, s, val, x + 1);
  54       } else {
  55         currentNode.token = s;
  56         currentNode.val = val;
  57         return currentNode;
  58       }
  59     } else if (currentNode.splitchar > s.charAt(x)) {
  60       currentNode.loKid = insert(currentNode.loKid, s, val, x);
  61     } else if (currentNode.splitchar == s.charAt(x)) {
  62       if (x < s.length() - 1) {
  63         currentNode.eqKid = insert(currentNode.eqKid, s, val, x + 1);
  64       } else {
  65         currentNode.token = s;
  66         currentNode.val = val;
  67         return currentNode;
  68       }
  69     } else {
  70       currentNode.hiKid = insert(currentNode.hiKid, s, val, x);
  71     }
  72     return currentNode;
  73   }
  74
  75   /**
  76    * Auto-completes a given prefix query using Depth-First Search with the end
  77    * of prefix as source node each time finding a new leaf to get a complete key
  78    * to be added in the suggest list.
  79    *
  80    * @param root
  81    *          a reference to root node of TST.
  82    * @param s
  83    *          prefix query to be auto-completed.
  84    * @param x
  85    *          index of current character to be searched while traversing through
  86    *          the prefix in TST.
  87    * @return suggest list of auto-completed keys for the given prefix query.
  88    */
  89   public ArrayList<TernaryTreeNode> prefixCompletion(TernaryTreeNode root,
  90           String s, int x) {
  91
  92     TernaryTreeNode p = root;
  93     ArrayList<TernaryTreeNode> suggest = new ArrayList<TernaryTreeNode>();
  94
  95     while (p != null) {
  96       if (s.charAt(x) < p.splitchar) {
  97         p = p.loKid;
  98       } else if (s.charAt(x) == p.splitchar) {
  99         if (x == s.length() - 1) {
 100           break;
 101         } else {
 102           x++;
 103         }
 104         p = p.eqKid;
 105       } else {
 106         p = p.hiKid;
 107       }
 108     }
 109
 110     if (p == null) return suggest;
 111     if (p.eqKid == null && p.token == null) return suggest;
 112     if (p.eqKid == null && p.token != null) {
 113       suggest.add(p);
 114       return suggest;
 115     }
 116
 117     if (p.token != null) {
 118       suggest.add(p);
 119     }
 120     p = p.eqKid;
 121
 122     Stack<TernaryTreeNode> st = new Stack<TernaryTreeNode>();
 123     st.push(p);
 124     while (!st.empty()) {
 125       TernaryTreeNode top = st.peek();
 126       st.pop();
 127       if (top.token != null) {
 128         suggest.add(top);
 129       }
 130       if (top.eqKid != null) {
 131         st.push(top.eqKid);
 132       }
 133       if (top.loKid != null) {
 134         st.push(top.loKid);
 135       }
 136       if (top.hiKid != null) {
 137         st.push(top.hiKid);
 138       }
 139     }
 140     return suggest;
 141   }
 142 }