19585번: 전설

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

 

언어 : Kotlin

 

문제 설명 :

Sogang ICPC Team에는 색상 이름과 닉네임의 순서로 이여서 팀명을 지으면 ICPC 리저널에서 수상할 수 있다는 전설이 있다. 색상 이름들과 닉네임들이 주어질 때, Q개의 팀에 대해 다음 리저널에서 수상할 수 있을지 전설에 기반해 알려주는 프로그램을 작성하자.

 

입력 :

첫째 줄에는 색상의 종류 C와 닉네임의 개수 N이 주어진다. (1 ≤ C, N ≤ 4,000)
다음 C개의 줄에는 색상 이름 C개가 한 줄에 하나씩 주어진다. 색상 이름은 중복되지 않는다.
다음 N개의 줄에는 Sogang ICPC Team 구성원들의 닉네임 N개가 한 줄에 하나씩 주어진다. 닉네임도 중복되지 않는다.
다음 줄에는 팀의 개수 Q가 주어진다. (1 ≤ Q ≤ 20,000)
다음 Q개의 줄에는 팀명 Q개가 한 줄에 하나씩 주어진다. 팀명은 중복되지 않는다.
모든 색상 이름의 길이와 닉네임의 길이는 1,000글자를 넘지 않는다. 모든 팀명은 2,000글자를 넘지 않는다. 모든 문자열은 알파벳 소문자로만 이루어져 있다.

 

출력 :

팀명이 입력된 순서대로, 전설에 따라 해당 팀이 다음 리저널에서 수상할 수 있다면 Yes, 아니라면 No를 출력한다.

 

제한 사항 :

• 시간 제한 : 3초
• 메모리 제한 : 1024MB

 

입출력 예 :

입력	출력
4 3 red blue purple orange shift joker noon 5 redshift bluejoker purplenoon orangeshift whiteblue	Yes Yes Yes Yes No

 

풀이 : 

무작정 풀었다가 시간 초과 당하고 트라이 + 셋으로 바꿔 풀었다.

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

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

    val (c, n) = readLine().split(" ").map { it.toInt() }
    val color = Trie().apply {
        for (i in 0 until c) {
            insert(readLine())
        }
    }

    val name: Set<String> = buildSet {
        for (i in 0 until n) {
            add(readLine())
        }
    }

    for (i in 0 until readLine().toInt()) {
        val cur = readLine().trim()
        bw.appendLine(if (color.search(cur, color, name)) "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) {
            cur = cur.children[i - 'a'] ?: run {
                val next = TrieNode()
                cur.children[i - 'a'] = next
                next
            }
        }

        cur.isLast = true
    }

    fun search(key: String, color: Trie, name: Set<String>): Boolean {
        var cur = color.root

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

            if (cur.isLast && key.substring(i) in name) return true
            if (cur.children[temp - 'a'] == null) break

            cur = cur.children[temp - 'a']!!
        }

        return false
    }
}

class TrieNode {
    var isLast: Boolean = false
    val children: Array<TrieNode?> = Array(26) { null }
}