1787번: 문자열의 주기 예측

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

 

언어 : Kotlin

 

문제 설명 :

알파벳 소문자들로만 이루어진 문자열을 생각하자. 이런 문자열을 읽어 나가다 보면, 문자열의 주기가 예측되는 순간이 있다. 다음과 같은 문자열을 예로 들어 보자.

a b a b a b a

이 문자열을 네 번째 문자까지의 문자열 'a b a b'와, 그 뒤에 남은 'a b a'로 나누어 생각해 볼 수 있다. 이렇게 하면 뒤쪽 문자열은 앞쪽 네 개의 문자 중 세 번째 문자까지가 반복되다가 끝나는 꼴이다.

a b a b a b a

또한, 여섯 번째 문자까지의 문자열 'a b a b a b'와, 그 뒤에 남은 'a'로 나누어서 생각할 수도 있다. 이 경우에도 뒤쪽 문자열은 앞쪽 문자열이 반복되다가 끝나는 꼴이다.

즉, 예시된 문자열은 'a b a b'혹은, 'a b a b a b'가 반복되는 문자열의 일부라고 예상할 수 있는 것이다. 단, 이러한 추정에서 뒤쪽 문자열이 앞쪽 문자열보다 길면 안 된다. 예를 들어 'a b'는 이 문자열의 주기로 예측하기에는 너무 짧다.

이제는 어떤 문자열 S에 대해, 첫 번째 문자부터, i번째 문자까지로 이루어진 부분 문자열 Si를 생각해 보자. 모든  각각에 대해, 위에 예시된 것처럼 문자열의 주기를 추정할 수 있다. 우리가 관심 있는 것은 이렇게  각각에 Si에 대해 추정할 수 있는 주기 중 가장 긴 것의 길이이다. 이를 Pi라고 하자. 예시된 문자열에서 P7은 4와 6중 최댓값인 6이 된다.

길이가 n인 문자열 S가 주어졌을 때, P1 + P2 + ... + Pn을 구하는 프로그램을 작성하시오.

 

입력 :

첫째 줄에 문자열 S의 길이 n이 주어진다. (1 ≤ n ≤ 1,000,000) 둘째 줄에 문자열 S가 공백 없이 주어진다.

 

출력 :

첫째 줄에, P1 + P2 + ... + Pn의 값을 출력한다.

 

제한 사항 :

• 시간 제한 : 2초
• 메모리 제한 : 128MB

 

입출력 예 :

입력	출력
8 babababa	24

 

풀이 : 

prefix와 suffix가 같은 가장 짧은 길이를 구하는 문제로, KMP 알고리즘의 일부인 LPS를 얻는 함수와, DP, 메모이제이션을 사용해 문제를 해결한다.

 

KMP의 실패 함수에 대한 이해 및 메모이제이션 최적화에 대한 문제.

실패 함수로 인해 아호코라식 알고리즘을 떠올릴 수도 있다.

 

DP배열과 메모이제이션으로 DP에 최솟값(가장 짧은 길이)을 저장해내며, 메모이제이션을 통해 이미 계산된 값을 재사용하는 방식으로 최적화를 이룬다.

import java.io.BufferedWriter
import java.io.OutputStreamWriter

const val MAX_VALUE = Long.MAX_VALUE

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

    val n = readLine().toInt()
    val s = readLine()
    val pi = getPi(s)

    var result = 0L
    for (i in 0 until n) {
        result += pi[i]
    }

    val dp = LongArray(n) { -1 }
    var sum = 0L

    for (i in 0 until n) {
        val t = computeMinLength(i, pi, dp)
        if (t < MAX_VALUE) {
            sum += (i - t + 1)
        }
    }

    bw.write("$sum")
    bw.flush()
    bw.close()
}

fun getPi(s: String): IntArray {
    val length = s.length
    val lps = IntArray(length)
    var j = 0
    for(i in 1 until length) {
        while(j > 0 && s[i] != s[j]) j = lps[j - 1]
        if(s[i] == s[j]) lps[i] = ++j
    }
    return lps
}

fun computeMinLength(x: Int, pi: IntArray, dp: LongArray): Long {
    if (x < 0) return MAX_VALUE
    if (dp[x] != -1L) return dp[x]
    if (pi[x] == 0) return MAX_VALUE

    val result = minOf(pi[x].toLong(), computeMinLength(pi[x] - 1, pi, dp))
    dp[x] = result

    return result
}