[BOJ] Sum and Product - P4

TaeGN·2024년 9월 24일

algorithm kotlin

BOJ Platinum Challenge

목록 보기

94/114

문제풀이

각 지점에서 1이 나오는 개수의 누적합을 구한다.
1이 아닌 지점의 인덱스 리스트를 만든다.
1이 아닌 지점들의 구간을 순회하며 구간의 왼쪽 바깥의 연속하는 1의 개수, 구간의 오른쪽 바깥의 연속하는 1의 개수를 고려하여 mul == sum이 되는 개수를 카운팅한다.

주의사항

소요시간

50분

package 백준.Platinum.P4.p18207_SumAndProduct

fun main() {
    val N = readln().toInt()
    val nArr = readln().trim().split(" ").map(String::toInt)
    val oneCount = IntArray(N)
    val notOneIdxList = mutableListOf<Int>()
    for ((idx, n) in nArr.withIndex()) {
        oneCount[idx] = oneCount.getOrElse(idx - 1) { 0 } + if (n == 1) 1 else 0
        if (n > 1) notOneIdxList.add(idx)
    }
    val totalOneCount = oneCount.last()
    var result = 0L
    for (i in notOneIdxList.indices) {
        val l = notOneIdxList[i]
        var sum = nArr[l]
        var mul = nArr[l].toLong()
        for (j in (i + 1) until notOneIdxList.size) {
            val r = notOneIdxList[j]
            sum += nArr[r]
            mul *= nArr[r]
            if (mul - sum > totalOneCount) break
            val insideOneCount = oneCount[r] - oneCount[l]
            val diff = (mul - (sum + insideOneCount)).toInt()
            val leftOneCount = oneCount[l] - oneCount.getOrElse(notOneIdxList.getOrElse(i - 1) { -1 }) { 0 }
            val rightOneCount =
                oneCount.getOrElse(notOneIdxList.getOrElse(j + 1) { -1 }) { totalOneCount } - oneCount[r]
            val outsideOneCount = leftOneCount + rightOneCount
            result += maxOf(0, minOf(leftOneCount, rightOneCount, diff, outsideOneCount - diff) + 1)
        }
    }
    println(result)
}