package udo.logic;

import java.util.ArrayList;
import java.util.List;

public class TernarySearchTree {
    protected static class Node {
        private char myChar;
        private Node left, right, center;
        protected Boolean isWordEnd;
        
        public Node(char myChar, Boolean isWordEnd) {
            this.myChar = myChar;
            this.isWordEnd = isWordEnd;
        }
    }
    
    private Node root = null;
    
    private Node add(String s, Integer index, Node node) {
        if (node == null) {
            node = new Node(s.charAt(index), false);
        }
        
        if (s.charAt(index) < node.myChar) {
            node.left = add(s, index, node.left);
        } else if (s.charAt(index) > node.myChar) {
            node.right = add(s, index, node.right);
        } else {
            if (index + 1 == s.length()) {
                node.isWordEnd = true;
            } else {
                node.center = add(s, index + 1, node.center); 
            }
        }
        
        return node;
    }
    
    public void add(String s) {
        if (s == null) {
            return;
        }
        
        s = s.trim();
        if (s.equalsIgnoreCase("")) {
            return;
        }
        
        root = add(s, 0, root);
    }
    
    public List<String> searchPrefix(String prefix, Integer numWords) {
        int index = 0;
        Node curNode = root;
        
        while (curNode != null) {
            if (index == prefix.length()) {
                return searchTree(curNode.center, numWords);
            }
            
            if (prefix.charAt(index) < curNode.myChar) {
                curNode = curNode.left;
            } else if (prefix.charAt(index) > curNode.myChar) {
                curNode = curNode.right;
            } else {
                index++;
                curNode = curNode.center;
            }
        }
        
        return null;
    }
    
    public List<String> searchPrefix(String prefix) {
        return searchPrefix(prefix, null);
    }
    
    /**
     * Find words in the tree rooted at some node
     * @param node the root of the tree for the search
     * @param numWords maximum number of words to search for or
     *                 search for all words if this is null
     * @return list of found words
     */
    private List<String> searchTree(Node node, Integer numWords) {
        List<Character> curWord = new ArrayList<>();
        List<String> result = new ArrayList<>();
        
        searchTreeHelper(node, numWords, curWord, result);
        
        return result;
    }

    private void searchTreeHelper(Node node, Integer numWords,
                                  List<Character> curWord,
                                  List<String> result) {
        if (numWords != null && result.size() == numWords) {
            return;
        }
        if (node == null) {
            return;
        }
        
        searchTreeHelper(node.left, numWords, curWord, result);
        if (numWords != null && result.size() == numWords) {
            return;
        }

        curWord.add(node.myChar);

        if (node.isWordEnd) {
            result.add(charListToString(curWord));
        }
        searchTreeHelper(node.center, numWords, curWord, result);
        
        curWord.remove(curWord.size() - 1);

        searchTreeHelper(node.right, numWords, curWord, result);
    }

    private String charListToString(List<Character> curWord) {
        StringBuilder builder = new StringBuilder(curWord.size());

        for (char c : curWord) {
            builder.append(c);
        }
        
        return builder.toString();
    }

    public static void main(String[] args) {
        
    }

}

@@ -1448,140 +1448,8 @@
 ) %7B%0A
-            if (index == prefix.length()) %7B%0A                return searchTree(curNode.center, numWords);%0A            %7D%0A            %0A
     
@@ -1692,16 +1692,345 @@
 ndex++;%0A
+                if (index == prefix.length()) %7B%0A                    if (curNode.isWordEnd) %7B%0A                        return searchTree(prefix, true, curNode.center, numWords);%0A                    %7D else %7B%0A                        return searchTree(prefix, false, curNode.center, numWords);%0A                    %7D%0A                %7D%0A
         
@@ -2293,16 +2293,167 @@
 ome node
+ appended by prefix%0A     * @param prefix the prefix from the root to this tree%0A     * @param prefixIncluded indicate whether the prefix is a valid word
 %0A     * 
@@ -2703,16 +2703,91 @@
 rchTree(
+String prefix, Boolean prefixIncluded,%0A                                    
 Node nod
@@ -2843,32 +2843,39 @@
 rWord = 
-new ArrayList%3C%3E(
+stringToCharList(prefix
 );%0A     
@@ -2939,58 +2939,440 @@
     
-searchTreeHelper(node, numWords, curWord, result);
+if (prefixIncluded) %7B%0A            result.add(prefix);%0A        %7D%0A        searchTreeHelper(node, numWords, curWord, result);%0A        %0A        return result;%0A    %7D%0A%0A    private List%3CCharacter%3E stringToCharList(String prefix) %7B%0A        List%3CCharacter%3E result = new ArrayList%3C%3E();%0A%0A        if (prefix != null) %7B%0A            for (int i = 0; i %3C prefix.length(); i++) %7B%0A                result.add(prefix.charAt(i));%0A            %7D%0A        %7D
 %0A   
@@ -4523,22 +4523,388 @@
 ) %7B%0A        
+TernarySearchTree t = new TernarySearchTree();%0A        %0A        t.add(%22cat%22); t.add(%22category%22); t.add(%22catalyzt%22);%0A        t.add(%22dog%22); t.add(%22dogmatic%22); t.add(%22dogwood%22);%0A        t.add(%22add%22); t.add(%22addition%22); t.add(%22additional%22);%0A        %0A        List%3CString%3E r = t.searchPrefix(%22ad%22);%0A        for (String s : r) %7B%0A            System.out.println(s);%0A        %7D
 %0A    %7D%0A%0A%7D%0A

Test and fix some bugs in ternary search tree