백준 1292: 쉽게 푸는 문제

solved.ac 실버 5에 랭크되어 있지만 재미있는 문제다.

어지간히 복잡한 방법을 사용해서 해결하더라도 수열의 길이가 1,000을 넘지 않는다는 점에서 해결이 어려운 문제는 아니다.

하지만 상수 시간에 해결할 수 있다면 어떨까?

 

문제에서 주어진 수열은 다음과 같다:

\(a_n\) = 1 2 2 3 3 3 4 4 4 4 5 5 5 ...

1이 1번, 2가 2번, 3이 3번, 4가 4번 나오는 방법이 반복되는 것이다.

이때 수열 \(a_n\)의 \(n\)번째 항까지의 합을 \(S(n)\)라고 하자.

문제에서 입력 A와 B가 주어질 때 A부터 B번째 수까지의 합은 다음과 같이 표현할 수 있다.

$$S(B) - S(A) + a_A$$

그렇다면 \(S(n)\)와 A번째 항 \(a_A\)를 상수 시간에 구할 수 있을까?

 

1, 2, 3, 4, 5 ... 꼴의 등차수열을 살펴보자.

n번째 항까지의 합은 초등학교 때의 가우스덕분에 상수시간에 해결할 수 있음을 알 수 있다.

$$\sum_{i=1}^{n}i={n\times(n+1)\over{2}}$$

그렇다면 어떤 정수 n이 주어졌을 때 n에 가장 가깝지만 크지 않은 수열의 합을 이루는 마지막 원소도 구할 수 있다.

$$\left \lfloor {\sqrt{8n + 1} - 1} \over 2 \right \rfloor$$

이를 이용하여 \(a_n\)은 어렵지 않게 구할 수 있게 되었다. 위 식에서 구한 마지막 원소를 \(l\), \(l\)까지의 자연수의 합을 \(s\)라고 할 때,

$$nth(n, l, s) = \begin{cases} l & \text{ if } n = s \\ l + 1 & \text{ if } n \neq s \end{cases}$$

 

또, \(S(n)\)는 이렇게 표현할 수 있다.

$$S(n) = \sum_{i=1}^{n}i^2 + (n-s) \times nth(n, l, s)$$

보통 급수 기호가 나온다면 반복을 떠올릴 수 있지만, 고등학교 이상의 수학을 배웠다면 저 급수 또한 간단한 식으로 바꿀 수 있다는 것을 알 수 있다.

$$\sum_{i=1}^{n}i^2 = {n \times (n + 1) \times (2n + 1) \over 6}$$

 

결국 반복 한 번 돌지 않고 해결할 수 있는 문제였음을 알 수 있다.

다음은 Kotlin으로 작성한 예제이다:

import java.io.BufferedReader
import java.io.BufferedWriter
import java.io.InputStreamReader
import java.io.OutputStreamWriter
import kotlin.math.floor
import kotlin.math.sqrt

// 초항이 1이고 각 항이 i의 제곱인 수열의 n번째 항까지의 합
// 1 4 9 16 25 36 49 ...
// ex n = 3) 14
private fun sigmaSquare(n: Int): Int {
    return (n * (n + 1) * (2 * n + 1)) / 6
}

// 초항이 1이고 공차가 1인 등차수열의 n번째 항까지의 합
private fun naturalSum(n: Int): Int {
    return (n * (1 + n)) / 2
}

// 초항이 1이고 공차가 1인 등차수열에서 n이 주어질 때 n보다 작거나 같은 최대 합 sum과 그 수 last를 반환
// ex n = 12) last = 4, sum = 10
private fun maxNaturalSum(n: Int): Pair<Int, Int> {
    val last = ((sqrt((8 * n.toLong() + 1).toDouble()) - 1) / 2).toInt()
    val sum = naturalSum(last)

    return last to sum
}

// 문제에서 주어진 수열에서 n번째 수를 반환
// 1 2 2 3 3 3 4 4 4 4 5 ...
// ex n = 10) 4
//    n = 11) 5
private fun nthNumber(n: Int, last: Int, sum: Int): Int {
    if (sum == n) {
        return last
    } else {
        return last + 1
    }
}

// 문제에서 주어진 수열에서 n번째 항까지의 합을 반환
// 1 2 2 3 3 3 4 4 4 4 5 ...
// ex n = 10) 30
//    n = 11) 35
private fun sumUntil(n: Int, last: Int, sum: Int): Int {
    return sigmaSquare(last) + (n - sum) * (last + 1)
}

fun main() {
    val writer = BufferedWriter(OutputStreamWriter(System.out))
    val reader = BufferedReader(InputStreamReader(System.`in`))
    val (a, b) = reader.readLine().split(" ").map(String::toInt)
    val (lastA, sumA) = maxNaturalSum(a)
    val (lastB, sumB) = maxNaturalSum(b)
    val sA = sumUntil(a, lastA, sumA)
    val sB = sumUntil(b, lastB, sumB)
    val answer = sB - sA + nthNumber(a, lastA, sumA)

    writer.write("$answer\n")
    writer.flush()
}