[백준] 9250번 : 문자열 집합 판별

https://www.acmicpc.net/problem/9250

 

해당 문제는 N개의 문자들이 Q번 입력받은 문자열들의 부분집합으로 포함되어 있는지 확인하는 문제입니다.
Trie와 KMP를 활용한 aho-corasick이라는 알고리즘 통해 해결항 수 있습니다.

 

• 트리에 insert

void insert(String word) {
		Trie curNode = this;
		for(int i = 0; i < word.length(); i++) {
			char c = word.charAt(i);
			
			curNode.child.putIfAbsent(c, new Trie());
			curNode = curNode.child.get(c);
			
			if(i == word.length() - 1) {
				curNode.output = true;
			}
		}
	}

 

• kmp의 fail함수를 트리형태로

public void computeFailFunc() {
		Queue<Trie> queue = new LinkedList<>();
		this.fail = this;
		queue.add(this);
		
		while(!queue.isEmpty()) {
			Trie cur = queue.poll();
			for(int i = 0; i < 26; i++) { // 모든 알파벳이 포함된 문자열 탐색이니 a ~ z
				char c = (char)(i + 97); 
				
				Trie nxt = cur.child.get(c);
				if(nxt == null) continue;
				
				if(cur == this) { // 1레벨 노드인 경우
					nxt.fail = this;
				}else {  
					Trie failLinkNode = cur.fail; // 부모의 실패 노드
					while(failLinkNode != this && failLinkNode.child.get(c) == null) {
						failLinkNode = failLinkNode.fail;
					} // 부모의 실패 노드의 자식 노드가 c가 아닌 경우 거슬러 탐색
					
					if(failLinkNode.child.get(c) != null) {
						failLinkNode = failLinkNode.child.get(c);
					} // 부모의 실패 노드의 자식 노드가 c인 경우 연결
					
					nxt.fail = failLinkNode;
				}
				
				if(nxt.fail.output) {
					nxt.output = true;
				}
				queue.add(nxt);
			}
		}
	}

 

• 검색

public boolean ahoCorasick(String word) {
		Trie curNode = this;
		for(int i = 0; i < word.length(); i++) {
			char c = word.charAt(i);
			while(curNode != this && curNode.child.get(c) ==null) {
				curNode = curNode.fail;
			}
			if(curNode.child.get(c)!=null) {
				curNode = curNode.child.get(c);
			}
			
			if(curNode.output) {
				return true;
			}
		}
		return false;
	}

 

• 전체코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;
import java.util.Queue;

public class Main {
	public static void main(String[] args) throws IOException{ 
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		
		int n = Integer.parseInt(br.readLine());
		
		Trie trieSet = new Trie();
		for(int i = 0; i < n; i++) {
			trieSet.insert(br.readLine());
		}
		
		trieSet.computeFailFunc();
		
		int q = Integer.parseInt(br.readLine());
		
		StringBuilder sb = new StringBuilder();
		for(int i = 0; i < q; i++) {
			sb.append(trieSet.ahoCorasick(br.readLine()) ? "YES" : "NO").append("\n");
		}
		System.out.print(sb);
	}
}

class Trie {
	Map<Character, Trie> child = new HashMap<>();
	boolean output;
	Trie fail;
	
	void insert(String word) {
		Trie curNode = this;
		for(int i = 0; i < word.length(); i++) {
			char c = word.charAt(i);
			
			curNode.child.putIfAbsent(c, new Trie());
			curNode = curNode.child.get(c);
			
			if(i == word.length() - 1) {
				curNode.output = true;
			}
		}
	}
	
	public void computeFailFunc() {
		Queue<Trie> queue = new LinkedList<>();
		this.fail = this;
		queue.add(this);
		
		while(!queue.isEmpty()) {
			Trie cur = queue.poll();
			for(int i = 0; i < 26; i++) { // 모든 알파벳이 포함된 문자열 탐색이니 a ~ z
				char c = (char)(i + 97); 
				
				Trie nxt = cur.child.get(c);
				if(nxt == null) continue;
				
				if(cur == this) { // 1레벨 노드인 경우
					nxt.fail = this;
				}else {  
					Trie failLinkNode = cur.fail; // 부모의 실패 노드
					while(failLinkNode != this && failLinkNode.child.get(c) == null) {
						failLinkNode = failLinkNode.fail;
					} // 부모의 실패 노드의 자식 노드가 c가 아닌 경우 거슬러 탐색
					
					if(failLinkNode.child.get(c) != null) {
						failLinkNode = failLinkNode.child.get(c);
					} // 부모의 실패 노드의 자식 노드가 c인 경우 연결
					
					nxt.fail = failLinkNode;
				}
				
				if(nxt.fail.output) {
					nxt.output = true;
				}
				queue.add(nxt);
			}
		}
	}
	
	public boolean ahoCorasick(String word) {
		Trie curNode = this;
		for(int i = 0; i < word.length(); i++) {
			char c = word.charAt(i);
			while(curNode != this && curNode.child.get(c) ==null) {
				curNode = curNode.fail;
			}
			if(curNode.child.get(c)!=null) {
				curNode = curNode.child.get(c);
			}
			
			if(curNode.output) {
				return true;
			}
		}
		return false;
	}
}