AC解 - 统计带某个前缀的单词数量

原题：http://acm.hdu.edu.cn/showproblem.php?pid=1251

Problem Description

Ignatius最近遇到一个难题,老师交给他很多单词(只有小写字母组成,不会有重复的单词出现),现在老师要他统计出以某个字符串为前缀的单词数量(单词本身也是自己的前缀).

 

Input

输入数据的第一部分是一张单词表,每行一个单词,单词的长度不超过10,它们代表的是老师交给Ignatius统计的单词,一个空行代表单词表的结束.第二部分是一连串的提问,每行一个提问,每个提问都是一个字符串.

注意:本题只有一组测试数据,处理到文件结束.

 

Output

对于每个提问,给出以该字符串为前缀的单词的数量.

 

Sample Input

banana
band
bee
absolute
acm

ba
b
band
abc

Sample Output

2
3
1
0

分析：利用Patricia前缀树（参考另一篇博文），再加一个统计叶子结点（带#终结符）的个数的函数getLeafCount()。



代码：
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
 

public class Main {
	private class PatriciaTrieNode {
	    private String key;
	    private Integer data;
	    private List<PatriciaTrieNode> children = new LinkedList<PatriciaTrieNode>();
	    
	    //use "#" for terminal char
	    private boolean terminal; 
	    
	    public PatriciaTrieNode(){	    	
	    	this.key = "";
	    }
	    public PatriciaTrieNode(String key){
	    	this.key = key;	    	
	    }	    
	    public String toString(){	    	
	    	return this.key + (this.terminal?"#":"") + "(" + children.size() +")";
	    }
	}
	private PatriciaTrieNode root;
	//value is only located at the external node
	private void insert(PatriciaTrieNode currNode,String key,Integer value) throws Exception{		
		boolean done = false;
		for(int i=0;i<currNode.children.size();i++){
			PatriciaTrieNode child = currNode.children.get(i);
			
			//use min(child.key.length, key.length)
			int len = child.key.length()<key.length()?child.key.length():
				key.length();
			int j = 0;
			for(;j<len;j++){
				if(key.charAt(j) != child.key.charAt(j)){
					break;
				}
			}
			if(j==0){//this child doesn't match	any character with the new key			
				//order keys by lexi-order
				if(key.charAt(0)<child.key.charAt(0)){
					//e.g. child="e" (currNode="abc")
					//	   abc                     abc
					//    /  \    =========>      / | \
					//   e    f   insert "c"     c# e  f
				
					PatriciaTrieNode node = new PatriciaTrieNode(key);
					currNode.children.add(i,node);
					node.terminal = true;		
					node.data = value;
					done = true;
					break;					
				}else{ //key.charAt(0)>child.key.charAt(0)
					//don't forget to add the largest new key after iterating all children
					continue;
				}
			}else{//current child's key partially matches with the new key; 0<j<=len				
				if(j==len){
					if(key.length()==child.key.length()){
						if(child.terminal){
							throw new Exception("Duplicate Key is found when insertion!");							
						}else{
							//e.g. child="ab"
							//	   ab                    ab#
							//    /  \    =========>    /   \
							//   e    f   insert "ab"  e     f
							child.terminal = true;
							child.data = value;
						}
					}else if(key.length()>child.key.length()){
						//e.g. child="ab#"
						//	   ab#                    ab#
						//    /  \    ==========>    / | \ 							
						//   e    f   insert "abc"  c# e  f						
						String subkey = key.substring(j);
						//recursion
						insert(child,subkey,value);
					}else{ //key.length()<child.key.length()
						//e.g. child="abc#"
						//	   abc#                      ab#
						//    /   \      =========>      /   
						//   e     f     insert "ab"    c#    
						//					           /  \
						//                            e    f													
						String childSubkey = child.key.substring(j); //c
						PatriciaTrieNode subChildNode = new PatriciaTrieNode(childSubkey);
						subChildNode.terminal = child.terminal;
						subChildNode.data = child.data;
						subChildNode.children = child.children; //inherited from parent
						
						child.key = key;  //ab
						child.terminal = true;  //ab#	
						child.data = value;
						
						child.children = new LinkedList<PatriciaTrieNode>();
						child.children.add(subChildNode);
					}					
				}else{//0<j<len
					//e.g. child="abc#"
					//	   abc#                     ab
					//    /  \     ==========>     / \
					//   e    f   insert "abd"    c#  d# 
					//                           /  \
					//                          e    f					
					//split at j
					String childSubkey = child.key.substring(j);  //c
					String subkey = key.substring(j); //d
					
					PatriciaTrieNode subChildNode = new PatriciaTrieNode(childSubkey);
					subChildNode.terminal = child.terminal;
					subChildNode.data = child.data;
					subChildNode.children = child.children; //inherited from parent
					
					//update child's key
					child.key = child.key.substring(0,j);
					//child is not terminal now due to split, it is inherited by subChildNode
					child.terminal = false;
					
					//Note: no need to merge subChildNode					
					
					PatriciaTrieNode node = new PatriciaTrieNode(subkey);
					node.terminal = true;
					node.data = value;
					child.children = new LinkedList<PatriciaTrieNode>();
					if(subkey.charAt(0)<childSubkey.charAt(0)){
						child.children.add(node);
						child.children.add(subChildNode);
					}else{
						child.children.add(subChildNode);
						child.children.add(node);
					}
				}
				done = true;
				break;
			}
		}
		if(!done){
			PatriciaTrieNode node = new PatriciaTrieNode(key);		
			node.terminal = true;
			node.data = value;
			currNode.children.add(node);
		}
	}
	public void insert(String key,Integer value) throws Exception{
		if(key == null || key.length() == 0) return;
		
		if(root==null){
			root = new PatriciaTrieNode();				
		}
		insert(root,key,value);		
	}
	
	public int query(String key){
		return query(root,key);
	}
	
	private int query(PatriciaTrieNode currNode,String key) {
		for(int i=0;i<currNode.children.size();i++){
			PatriciaTrieNode child = currNode.children.get(i);
			
			//use min(child.key.length, key.length)
			int len = child.key.length()<key.length()?child.key.length():
				key.length();
			int j = 0;
			for(;j<len;j++){
				if(key.charAt(j) != child.key.charAt(j)){
					break;
				}
			}
			
			if(j==0){//this child doesn't match	any character with the new key			
				//order keys by lexi-order
				if(key.charAt(0)<child.key.charAt(0)){
					//e.g. child="e", key="c" (currNode="abc")
					//	   abc                     
					//    /  \     
					//   e    h   
					return 0;
				}else{
					//e.g. child="e", key="h" (currNode="abc")
					continue;
				}
			}else{//current child's key partially matches with the new key; 0<j<=len				
				if(j==len){
					if(key.length()==child.key.length()){
						if(child.terminal){
							//e.g. child="ab", key="ab"
							//	   ab#                    
							//       \    
							//        f#    
							return getLeafCount(child);						
						}else{
							//e.g. child="ab", key="ab"
							//	   ab                    
							//    /  \    
							//   e    f    
							return getLeafCount(child);
						}
					}else if(key.length()>child.key.length()){
						//e.g. child="ab#", key="abc"
						//	   ab#                     
						//    /  \     						
						//   a    c#   			
						String subkey = key.substring(j); //c
						//recursion
						return query(child,subkey);
					}else{ //key.length()<child.key.length()
						//e.g. child="abc", key="ab"
						//	   abc                      
						//    /   \       
						//   e     f     
						return getLeafCount(child);						
					}					
				}else{//0<j<len
					//e.g. child="abc", key="abd"
					//	   abc                     
					//    /  \      
					//   e    f    
					return 0;					
				}				
			}
			
		}
		return 0;
	}
	
	public int getLeafCount(PatriciaTrieNode currNode){
		int count = 0;
		if(currNode.terminal)
			count++;
		for(int i=0;i<currNode.children.size();i++){
			PatriciaTrieNode child = currNode.children.get(i);
			count += getLeafCount(child);
		}
		return count;
	}
 
	
	public static void main(String[] args) {
		BufferedReader bf = new  BufferedReader(new InputStreamReader(System.in));
		List<String> queries = null;
		Main main = new Main();
		try{	
			String line = bf.readLine();
			line = line.trim();
			
			while(line.length()>0){
				try {
					main.insert(line, 0);
				} catch (Exception e) {
					
				}
				line = bf.readLine();
				line = line.trim();
			}
			//main.printTree();			
			
			queries = new ArrayList<String>();			
			while((line=bf.readLine())!= null){		
				line = line.trim();
				queries.add(line);
			}
		}catch(Exception e){
			
		}
		for(String str:queries){
			System.out.println(main.query(str));			
		}
	}

}