9250번: 문자열 집합 판별

문제 출처 : https://www.acmicpc.net/problem/9250

 

언어 : Kotlin

 

문제 설명 :

집합 S는 크기가 N이고, 원소가 문자열인 집합이다. Q개의 문자열이 주어졌을 때, 각 문자열의 부분 문자열이 집합 S에 있는지 판별하는 프로그램을 작성하시오. 문자열의 여러 부분 문자열 중 하나라도 집합 S에 있으면 'YES'를 출력하고, 아무것도 없으면 'NO'를 출력한다.

예를 들어, 집합 S = {"www","woo","jun"} 일 때, "myungwoo"의 부분 문자열인 "woo" 가 집합 S에 있으므로 답은 'YES'이고, "hongjun"의 부분 문자열 "jun"이 집합 S에 있으므로 답은 'YES'이다. 하지만, "dooho"는 모든 부분 문자열이 집합 S에 없기 때문에 답은 'NO'이다.

 

입력 :

첫째 줄에 집합 S의 크기 N이 주어진다. (1 ≤ N ≤ 1000)
다음 N개 줄에 집합 S의 원소들이 주어진다. 이 문자열의 길이는 100을 넘지 않는다.
다음 줄에 답을 판별해야 하는 문자열의 개수 Q가 주어진다. (1 ≤ Q ≤ 1000)
다음 Q개 줄에 답을 판별해야 하는 문자열이 주어진다. 이 문자열의 길이는 10000을 넘지 않는다.
입력으로 주어지는 모든 문자열은 알파벳 소문자로만 이루어져 있다.

 

출력 :

Q개 줄에 각 문자열에 대한 답을 출력한다.

 

제한 사항 :

• 시간 제한 : 1초
• 메모리 제한 : 256MB

 

입출력 예 :

입력	출력
3 www woo jun 3 myungwoo hongjun dooho	YES YES NO

 

풀이 : 

평범하게 중첩반복문 + contain으로 판별하려고 했다가 시간초과가 발생했다.

KMP로 풀어야하는 문제인가 싶었으나, 문제에서 주어진 최대치를 가정하고 생각해보면 그것도 아닌 것 같다는 생각이 들었고 확인해보니 KMP 알고리즘에 Trie 자료구조를 접목시킨 아호-코라(Aho-Corasick) 알고리즘을 사용해야한다는 것을 알게되었다.

 

아호-코라식 알고리즘은 다음 두 사이트를 참고하면 좋다.

https://loosie.tistory.com/587

[알고리즘] 다중 문자열 검색 아호 코라식(Aho-Corasick) 알고리즘 정리 (Java)

https://pangtrue.tistory.com/305

[알고리즘] 아호 코라식(Aho-Corasick) 알고리즘

 

아호-코라식 알고리즘 사용 시, O(N + M + P)의 시간복잡도로 해결이 가능하므로, 해당 알고리즘을 사용해야하는데 이 아호-코라식 알고리즘은 설명하기를 트라이 자료구조를 사용해 실패 링크와 출력 문자열 목록을 생성한 뒤, KMP 알고리즘과 같은 방식으로 문자열 매칭을 하면 된다는 것이다.

 

트라이 자료구조 방식에서 계산한 실패함수 자료구조

 

• 실선 화살표 = 해당 상태에서 대응이 성공했을 경우 이동가능한 상태
• 점선 화살표 = 실패 함수를 나타내는 것으로, KMP 알고리즘에선 getPi에 해당하는 부분이다.

 

구현 방식은 다음과 같았다.

1. Trie 자료 구조를 구현한다.
   1. 기본적으로 insert, find인데, find에 KMP 알고리즘을 구현, 추가로 실패링크 계산 로직인 failure 만든다.
2. 주어진 패턴을 전부 Trie에 insert한다.
3. 실패함수를 계산한다.
4. 주어진 문자열을 전부 Trie에 find하여 결과를 도출한다.

 

import java.io.BufferedReader
import java.io.BufferedWriter
import java.io.InputStreamReader
import java.io.OutputStreamWriter
import java.util.*


val trie = Trie()

fun main() = with(BufferedReader(InputStreamReader(System.`in`))) {
    val bw = BufferedWriter(OutputStreamWriter(System.out))

    val n = readLine().toInt()
    repeat(n) {
        trie.insert(readLine())
    }
    trie.failure()

    val q = readLine().toInt()
    repeat(q) {
        bw.appendLine(if (trie.find(readLine())) "YES" else "NO")
    }
    bw.flush()
    bw.close()
}

class Trie(private val root: TrieNode = TrieNode()) {
    fun insert(key: String) {
        var cur = root

        for (i in key.indices) {
            val next = key[i]

            cur.children.putIfAbsent(next, TrieNode())
            cur = cur.children[next]!!
        }

        cur.isFinish = true
    }

    fun failure() { // KMP 비교 중 실패했을 때 이동할 지점 정의하는 함수
        val queue: Queue<TrieNode> = LinkedList()
        root.failed = root
        queue.add(root)

        while (queue.isNotEmpty()) {
            val currentNode = queue.poll()

            for (i in 0 until 26) {
                val next = (97 + i).toChar()
                val nextNode = currentNode.children[next] ?: continue

                if (currentNode == root) {
                    nextNode.failed = root
                } else {
                    var failure = currentNode.failed

                    while (failure!! != root && failure.children[next] == null) {
                        failure = failure.failed
                    }

                    if (failure.children[next] != null) failure = failure.children[next]
                    nextNode.failed = failure
                }

                if (nextNode.failed!!.isFinish) nextNode.isFinish = true
                queue.add(nextNode)
            }
        }
    }

    fun find(key: String): Boolean { // KMP
        var cur = root

        for (i in key.indices) {
            val next = key[i]

            while (cur != root && cur.children[next] == null) {
                cur = cur.failed!!
            }

            if (cur.children[next] != null) cur = cur.children[next]!!
            if (cur.isFinish) return true
        }

        return false
    }
}

class TrieNode {
    val children = HashMap<Char, TrieNode>()
    var failed: TrieNode? = null
    var isFinish = false
}