색상과 닉네임이 각각 최대 4,000개까지 주어진다.
각 문자열의 길이는 최대 1,000까지 가능하다.
이어서 최대 20,000개의 팀명이 입력된다,
팀명의 최대 길이는 2,000이며,
각 팀명이 색상+닉네임 으로 구성되는지 확인을 해야된다.
문자 하나 하나 다 확인을 하는 경우를 생각해보자.
팀명이 20,000개 주어졌고, 모든 길이가 2,000이라고 하면,
20,000 * 2,000 => 40,000,000 이다.
여기에 색상 및 닉네임이 최대 8,000개 입력되고 길이가 1,000이라면
총 8,000,000 이다.
두 숫자가 곱해진다면 제한 시간 내에 푸는 것은 절대 불가능하다.
두 숫자가 더해지는 형태의 시간 복잡도를 유도해야한다.
이럴 때 활용할 수 있는 자료구조가 트라이이다.
(시간복잡도를 끌어올린 만큼 공간복잡도는 희생된다)
트라이를 활용하면 문자열 하나당 문자열 길이 만큼만 탐색하면 된다.
다른 여러 문자열과의 조합을 판단할 필요가 없다.
문자열을 탐색하다가 매칭되는 데이터가 없다면 곧바로 중도 종료할 수도 있다.
우선 트라이를 구현해보자.
노드로 사용될 클래스를 TrieNode
로 선언한다.
data class TrieNode(
var key: Char? = null,
val children: Array<TrieNode?> = Array('z'-'a' + 1) { null },
var isTail: Boolean = false
)
key
에는 현재 알파벳을 기록하고, children
에는 이어질 문자 TrieNode
를 기록한다.
children
의 인덱스는 문자의 ASCII 값을 활용한다.
(Array
대신 Map
을 사용할 수도 있으나, 본 문제에서는 시간 초과가 발생하는 관계로 배열을 활용한다.)
이제 Trie
클래스를 구현한다.
멤버 변수 root
가 최상위 TrieNode
역할을 한다.
문자열을 기록하는 insert
함수도 선언한다.
// 문자의 ASCII 를 배열의 인덱스로 사용한다.
fun Char.toIndex() = this - 'a'
class Trie {
val root = TrieNode() // 최상위 노드
// 문자열 기록 함수
fun insert(color: String) {
var ptr = root // 노드 포인터
for (c in color) { // 기록할 문자열의 문자 하나 하나를 기록한다.
// 현 노드의 children 을 가져와서 문자 c에 해당하는 TrieNode를 가져온다.
ptr = ptr.children[c.toIndex()] ?: run {
// ptr.children[c.toIdex()] == null 이라면, 등록해주고 리턴한다.
val nextNode = TrieNode(c)
ptr.children[c.toIndex()] = nextNode
nextNode
}
}
// 만약 삽입한 문자가 red라면
// d에 해당하는 TrieNode가 ptr로 설정되어 있다.
// isTail을 true로 기록함으로써 문자열의 끝을 기록한다.
ptr.isTail = true
}
}
클래스가 구현됐으니 입력을 받자
색상은 Trie
에 저장하고, 닉네임은 Set
에 저장할 것이다.
색상까지는 Trie
를 통해 확인하고, 이어지는 문자열은 Set
에 존재하는지 확인하면서
시간복잡도를 조금 더 개선할 수 있다.
private lateinit var colorTrie: Trie
private lateinit var names: Set<String>
fun main() = with(System.`in`.bufferedReader()) {
val (C, N) = readLine().split(" ").map(String::toInt)
colorTrie = Trie().apply {
for (i in 0 until C) {
insert(readLine())
}
}
names = buildSet {
for (i in 0 until N) {
add(readLine())
}
}
...
}
이제 팀명을 입력받으면서 답이 될 수 있는지 확인한다. search
함수를 선언할 것이다.
colorTrie
변수의 root
부터 문자를 탐색해나간다.
트라이에 일치하는 문자가 없으면 그대로 탐색을 중단한다.
만약 색상 문자에 일치하면서 이어지는 문자열이 닉네임 Set
에 존재한다면, 답이 될 수 있다.
fun search(cand: String): Boolean {
var nodePtr = colorTrie.root
for (i in cand.indices) {
val s = cand[i]
// 문자의 끝이고, 남은 문자가 닉네임이라면
if (nodePtr.isTail && cand.substring(i) in names) {
return true
}
// 찾으려는 문자가 없으면 중단한다
if (nodePtr.children[s.toIndex()] == null) {
break
}
// 다음 문자의 TrieNode로 포인터를 설정
nodePtr = nodePtr.children[s.toIndex()]!!
}
return false
}
main
함수에서 Q
개 만큼의 팀명을 입력받으며 search
함수로 답인지를 판단한다.
StringBuilder
에 답을 누적한 뒤, 마지막에 출력하면 정답처리.
...
val answer = StringBuilder()
for (i in 0 until readLine().toInt()) {
val cand = readLine().trim()
if (search(cand))
answer.append("Yes").append("\n")
else
answer.append("No").append("\n")
}
print(answer)
티어가 높은 문제인데도 트라이를 학습하기에 굉장히 좋은 문제인 것 같다.