[BOJ] 순서쌍의 곱의 합 in Kotlin

문제풀러 가기!

생각이 좀 필요한 문제였습니다. 그리고 타입에 대해 신경을 썼어야 했던 문제였습니다.

문제 풀이 아이디어는 생각보다 복잡하지는 않습니다. 그래도 무지성으로 코드를 작성하면 당연히 문제가 틀릴겁니다.

일단 문제 풀이 아이디어에 대해서 생각해보죠.

1. 문제의 조건 상 10만개의 데이터가 주어집니다. O(n^2)으로 풀 생각은 접으셔야 합니다.
2. 문제에서 2개의 쌍을 곱한 모든 숫자를 더해야 한다고 했습니다. 조합으로 문제 풀까? 라는 생각도 하실 수 있는데, 이것도 접으셔야 합니다. nC2를 생각하실텐 데, 이것도 시간초과납니다.

대신에 문제 아이디어는 생각보다 빠르게 찾을 수 있습니다. 이걸 숫자로 나타내는 것이 아니라 문자로 생각해보겠습니다. 인수분해? 이런 마인드로 보시면 될 것 같습니다. 예를 들어서 4개의 숫자가 주어졌다고 생각해볼게요.

$a, b, c, d$

가 주어졌다고 생각해보면, 우리는 다음의 합을 구해야 합니다.

$ab+ac+ad+bc+bd+cd$

이걸 괄호를 쳐볼게요.

$(ab+ac+ad)+(bc+bd)+cd$

이걸 좀 더 보기 좋게 하면

$(ab+ac+ad) +(bc+bd) +cd → a(b+c+d) + b(c+d) +cd$

입니다. 결국에는 평소와 같이 누적합을 진행하지만, 약간은 다른 것을 확인할 수 있습니다.

즉 다음과 같은 수식으로 진행하면 누적합을 온전히 활용하면서 시간초과가 되지 않는 코드를 작성할 수 있습니다.

$a(a+b+c+d-a) + b(a+b+c+d-a-b) + c(a+b+c+d-a-b-c)$

즉 입력값 n에 대해서 수열 a와 그 합인 수열 S가 있다면 다음과 같은 같습니다.

마지막 우측식으로 문제를 풀어주시면 되겠습니다.

물론 $a_i=0$이면 continue 해주시면 됩니다.

따라서 풀이 코드는 다음과 같습니다. 아 그리고 Long형 잊지마세요.

import java.io.BufferedReader
import java.io.InputStreamReader

fun main() = with(BufferedReader(InputStreamReader(System.`in`))) {
    val n = readLine().toInt()
    val arr = readLine().split(" ").map { it.toInt() }
    val cumArr = IntArray(n)
    cumArr[0] = arr[0]
    for (i in 1 until n) {
        cumArr[i] = cumArr[i-1] + arr[i]
    }
    var answer: Long = 0
    for (i in 0 until n-1) {
        answer += ((cumArr[n-1].toLong()-cumArr[i].toLong())*arr[i].toLong())
    }
    print(answer)
    close()
}