数据结构之字典树 Java实现

数据结构之字典树 Java实现

public class Tire {

	private Tire[] children;

	private boolean isLeaf;

	private long passCount;

	private long leafCount;

	public static final Tire INSTANCE = new Tire();

	private Tire() {
		children = new Tire[26];
		isLeaf = false;
	}

	private Tire[] getChildren() {
		return children;
	}

	private void incPassCount() {
		passCount++;
	}

	private void incLeafCount() {
		leafCount++;
	}

	private boolean noPass(Tire node) {
		return node.passCount == 0;
	}

	public void insert(String word) {
		if (word == null || word.length() == 0 || search(word)) return;
		Tire node = this;
		for (int i = 0; i < word.length(); i++) {
			int index = word.charAt(i) - 'a';
			Tire[] children = node.getChildren();
			if (children[index] == null) {
				children[index] = new Tire();
			}
			node = children[index];
			node.incPassCount();
		}
		node.isLeaf = true;
		node.incLeafCount();
	}

	public boolean search(String word) {
		if (word == null || word.length() == 0) return false;
		return dfsSearch(this, 0, word);
	}

	private boolean dfsSearch(Tire node, int index, String word) {
		if (index == word.length()) {
			return node.isLeaf;
		}
		int ch = word.charAt(index) - 'a';
		Tire[] children = node.getChildren();
		if (children[ch] != null && !noPass(children[ch]) && dfsSearch(children[ch], index + 1, word)) {
			return true;
		}
		return false;
	}

	public boolean preSearch(String word) {
		if (word == null || word.length() == 0) return false;
		Tire node = this;
		for (int i = 0; i < word.length(); i++) {
			int ch = word.charAt(i) - 'a';
			Tire[] children = node.getChildren();
			if (children[ch] == null || noPass(children[ch])) return false;
			node = children[ch];
		}
		return true;
	}

	private Tire preSearchNode(String word) {
		if (word == null || word.length() == 0) return null;
		Tire node = this;
		for (int i = 0; i < word.length(); i++) {
			int ch = word.charAt(i) - 'a';
			Tire[] children = node.getChildren();
			if (children[ch] == null || noPass(children[ch])) return null;
			node = children[ch];
		}
		return node;
	}

	public List<String> getAllPreWord(String word) {
		Tire tire = preSearchNode(word);
		if (tire == null || noPass(tire)) return new ArrayList<>();
		return dfsAllPreWord(tire, word);
	}

	public List<String> getAllWord(){
		return dfsAllPreWord(this,"");
	}

	private List<String> dfsAllPreWord(Tire node, String word) {
		List<String> res = new ArrayList<>();
		if (node.isLeaf) {
			if (node.passCount == 1) {
				return Arrays.asList(word);
			} else {
				res.addAll(Arrays.asList(word));
			}
		}
		Tire[] children = node.getChildren();
		for (int i = 0; i < children.length; i++) {
			if (children[i] != null && !noPass(children[i])) {
				char ch = (char) ('a' + i);
				res.addAll(dfsAllPreWord(children[i], word + ch));
			}
		}
		if (res.isEmpty()) {
			return Arrays.asList(word);
		}
		return res;
	}

	public void remove(String word) {
		if (search(word)) {
			dfsRemove(this, 0, word);
		}
	}

	private void dfsRemove(Tire node, int index, String word) {
		if (index == word.length()) {
			node.leafCount--;
			if (node.leafCount == 0) {
				node.isLeaf = false;
			}
			return;
		}
		int ch = word.charAt(index) - 'a';
		Tire[] children = node.getChildren();
		children[ch].passCount--;
		dfsRemove(children[ch], index + 1, word);
	}

	public static void main(String[] args) {
		Tire tire = new Tire();
		tire.insert("at");
		tire.insert("at");
		tire.insert("and");
		tire.insert("ac");
		tire.insert("ace");
		tire.insert("bc");
		tire.insert("an");
		tire.insert("add");
		tire.insert("adddd");
		tire.insert("bat");
		tire.insert("abc");
		tire.remove("at");
		System.out.println(tire.search("at"));
		System.out.println(tire.getAllPreWord("an"));
		System.out.println(tire.getAllWord());
	}
}