//R-ways TST
import edu.princeton.cs.algs4.Queue;

public class RTST {
    private static final int R = 256;
    private Node root;

    private class Node {
        Object val;
        Node[] next = new Node[R];
    }

    public Object get(String key) {
        Node x = get(root, key, 0);
        if (x == null) return null;
        else return x.val;
    }

    private Node get(Node x, String key, int d) {
        if (x == null) return null;
        if (d == key.length()) return x;
        get(x.next[key.charAt(d)], key, d + 1);
    }

    public void put(String key, Object val) {
        root = put(root, key, val, 0);
    }

    private Node put(Node x, String key, Object val, int d) {
        if (x == null) x = new Node();
        if (d == key.length()) {
            x.val = val;
            return x;
        }
        x.next[key.charAt(d)] = put(x.next[key.charAt(d)], key, val, d + 1);
        return x;
    }

    public Iterable<String> keys() {
        Queue<String> queue = new Queue<String>();
        collect(root, "", queue);
    }

    private void collect(Node x, String prefix, Queue<String> queue) {
        if (x == null) return;
        if (x.val != null) queue.enqueue(prefix);
        for (int i = 0; i < R; i++) {
            collect(x.next[i], prefix + i, queue);
        }
    }

    public Iterable<String> keysWithPrefix(String pre) {
        Queue<String> queue = new Queue<String>();
        Node x = get(root, pre, 0);
        collect(x, pre, queue);
        return queue;
    }

    public String longestPrefix(String s) {
        int length = search(root, s, 0, 0);
        return s.substring(0, length);
    }

    private int search(Node x, String s, int d, int length) {
        if (x == null) return length;
        if (x.val != null) length = d;
        if (d == s.length()) return length;
        return search(x.next[s.charAt(d)], s, d + 1, length);
    }

    public Iterable<String> keysThatMatch(String pat) {
        Queue<String> queue = new Queue<String>();
        collect(root, "", pat, queue);
        return queue;
    }

    private void collect(Node x, String prefix, String pat, Queue<String> queue) {
        int d = prefix.length();
        if (x == null) return;
        if (d == pat.length() && x.val != null) queue.enqueue(prefix);
        char next = pat.charAt(d);
        for (char c = 0; c < R; c++) {
            if (next == '.' || next == c) {
                collect(x.next[c], prefix + c, pat, queue);
            }
        }
    }

    public void delete(String key) {
        root = delete(root, key, 0);
    }

    private Node delete(Node x, String key, int d) {
        if (x == null) return null;
        if (d == key.length()) {
            x.val = null;
        } else {
            char c = key.charAt(d);
            x.next[c] = delete(x.next[c], key, d + 1);
        }

        if (x.val != null) return x;
        for (char c =0; c < R;c++) {
            if (x.next[c] != null return x;
        }

        return null;
    }
}

//3-ways TST
import edu.princeton.cs.algs4.Queue;

public class TST {
    private Node root = new Node();

    private class Node {
        Object val;
        char c;
        Node left,middle,right;
    }

    public Object get(String key) {
        Node x = get(root, key, 0);
        return x.val;
    }

    private Node get(Node x, String key, int d) {
        if (x == null) return null;
        char c = key.charAt(d);
        if (d == key.length() - 1) {
            if (x.val != null) return x;
            else  return null;
        }

        if (c < x.c) return get(x.left, key, d);
        if (c > x.c) return get(x.right, key, d);
        if (c == x.c) return get(x.middle, key, d+1);
    }

    public void put(String key, Object val) {
        root = put(root, key, val, 0);
    }

    private Node put(Node x, String key, Object val, int d) {
        char c = key.charAt(d);
        if (x == null) {
            x = new Node();
            x.c = c;
        }

        if (d == key.length() - 1) {
            x.val = val;
        }

        if (c < x.c) x.left = put(x.left, key, val, d);
        if (c > x.c) x.right =  put(x.right, key, val, d);
        if (c == x.c) x.middle = put(x.middle, key, val, d+1);

        return x;
    }

    public Iterable<String> keys() {
        Queue<String> queue = new Queue<>();
       collect(root,queue,new StringBuilder()); //由于要修改prefix的值,故使用stringbuilder
        return queue;
    }

    private void collect(Node x, Queue<String> queue, StringBuilder prefix) {
        if (x == null) return ;
        if (x.val != null) queue.enqueue(prefix.toString() + x.c); // 别忘了加 x.c
        collect(x.left,queue,prefix);
        collect(x.middle,queue,prefix.append(x.c));//
        prefix.deleteCharAt(prefix.length() - 1); // prefix的值在上一行中改变了
        collect(x.right,queue,prefix);
    }

    public Iterable<String> keysWithPrefix(String pre) {
        Queue<String> queue = new Queue<>();
        Node x = get(root, pre, 0);
        if (x.val != null) queue.enqueue(pre);
        collect(x,queue,new StringBuilder(pre));
        return queue;
    }

    public String longestPrefixOf(String query) {
        int length = search(root, query, 0, 0);
        return query.substring(0,length - 1);
    }

    private int search(Node x, String query, int d, int length) {
        char c = query.charAt(d);
        if (x == null) return length;
        if (x.val != null) length = d;
        if (d == query.length() - 1) return length;

        if (c < x.c) return search(x.left, query, d, length);
        else if(c > x.c) return search(x.right, query, d, length);
        else return search(x.middle, query, d + 1, length);
    }
}