백준 2632번: 피자 판매 (Kotlin)

https://www.acmicpc.net/problem/2632

2632번: 피자판매

안녕하세요, 이륙사입니다.

 

풀이

이번 문제는 특정 알고리즘을 사용해서 완전 탐색(brute force)을 하는 문제였습니다. 

 

저는 문제를 다음과 같이 해석했습니다. 

"배열A와 배열B의 숫자들을 더해서 특정 숫자를 만드는 경우의 수를 알고싶다. 단, 각 배열에서 숫자를 뽑을 땐 인덱스를 기준으로 연속해야 하며, 한쪽 배열의 숫자들만 사용할 수도 있다. 그리고 배열은 원형으로 되어있다."

 

즉, 각 원형 배열에서 길이 0 ~ (m or n) 으로 만들 수 있는 모든 연속 합을 각각 구하고, 그것을 A, B의 연속합 배열이라고 합시다. 각 연속 합 배열에서 숫자를 하나씩 골라서 더했을 때 원하는 숫자를 몇 개 만들 수 있는지 확인하면 정답을 구할 수 있습니다.

 

연속합 구하기

연속합을 구하는 과정은 다음과 같습니다.

길이1, 길이3일 때 연속 합

배열이 원형이기 때문에 인덱스를 넘어가더라도 시작하는 위치가 마지막 인덱스가 될 때까지 합을 구합니다.

ex) 길이 3 예시에서 빨간색 구간의 연속합: array[3] + array[4] + array[0] 

따라서 길이가 몇이든 항상 array의 전체 길이 만큼(위에선 5개) 연속합을 구하게 됩니다.

단, 한쪽 피자만 사용할 수도 있으며, 길이가 0이거나 전체인 경우에는 1개의 연속합이 나옵니다.

    var startSum = 0 // 첫 구간의 합
    for (sequentSize in 1 until pizza.size - 1) {
        startSum += pizza[sequentSize - 1] // 인덱스 0부터 연속 길이만큼의 피자 합

        var start = 0
        var end = sequentSize
        var sequentSum = startSum

        // 오른쪽으로 슬라이딩 하면서 길이 sequentSize에 해당하는 연속 합들을 구한다.
        // 마지막 인덱스가 첫번째 숫자가 될 때까지 진행한다.
        while(start <= pizza.lastIndex) {
            sumCount[sequentSum] += 1 // 해당 연속합이 몇 개 존재하는지 count
            sequentSum = sequentSum - pizza[start++] + pizza[end++] // 구간 오른쪽으로 이동
            if (end == pizza.size) {
                end = 0
            }
        }
    }
    sumCount[0] += 1 // 피자 0개 사용
    sumCount[startSum + pizza[pizza.lastIndex]] = 1 // 전체 합

 

누적 합 사용

위의 방법으로 연속 합을 구할 수 있습니다. 하지만 투포인터 형식을 사용해서 약간은 복잡해보입니다. 이때 누적 합 방식을 사용하면 좀 더 빠르고 편하게 구할 수 있습니다.

각 인덱스에서 시작해서 오른쪽으로 이동하면서 누적합을 구합니다. 그러면 각 인덱스를 시작으로 길이가 1 ~ n-1인 연속합을 모두 구할 수 있습니다.

for (i in 0 until pizza.size) { // 시작 위치
        var sum = 0

        // 항상 pizza-1개의 연속합을 찾을 수 있다(전체 합 제외)
        for (j in i until i + pizza.size - 1) {
            sum += pizza[j % pizza.size]
        }
    }

 

가지 수 계산

모든 연속합을 구했으니 이제 모두 확인해보면 답을 찾을 수 있습니다. 하지만 피자조각이 최대 1000개이기 때문에 전부 확인하면 시간초과가 발생합니다.

각 연속합을 구할 때 O(n^2), 두 연속합끼리 비교할 때 O(n^2) ==> 최종: O(n^4)

 

어떻게 시간 복잡도를 줄일 수 있을까요?

우리가 찾고자 하는 목표가 정해져있다는 것에 주목해봅시다. 연속합A에서 특정 합 k를 골랐을 때, 연속합B에서 target - k가 있는지 확인만 하면 됩니다. 즉, A에서 숫자를 골랐을 때 B에서는 그 숫자가 있는지 없는지, 있다면 몇 개가 있는지 바로 확인할 수 있다면 굳이 모든 조합을 비교하지 않아도 됩니다.

 

해싱

위의 방법대로 구현하기 위해서 Map(key: sum, value: count)을 사용할 수 있습니다. A의 연속합(sumA)을 구해서 모두 Map에 저장합니다. 그리고 B의 연속합(sumB)을 구해서 Map에 (target - sumB)이 있는지 즉, sumA + sumB = target을 만드는 sumA가 Map에 있는지 확인하면 O(n^2)의 시간으로 문제를 해결할 수 있습니다. 

// pizza A의 연속합을 구할 때마다 counArray에 sum count를 증가시킨다
// 숫자가 최대 1,000,000이기 때문에 map 대신 배열을 사용할 수 있음
sumCountArray[sequentSum] += 1


// pizza B의 연속합을 구할 때마다 target을 만드는 sumA가 있는지 확인
if (target - sum >= 0){
	count += sequentSumArray[target - sum] // 없으면 0, 있으면 1 이상이 저장되어 있음
}

연속합의 크기가 최대 1,000,000이기 때문에 Map 대신 배열을 사용했습니다.

 

투 포인터

원하는 숫자가 정해져있다는 것에서 힌트를 얻어 해싱 기법을 사용할 수 있었습니다. 마찬가지 이유로 투 포인터를 사용할 수 있습니다. 한쪽에서 연속합을 선택했을 때, 다른 쪽의 연속합이 정렬되어 있다면 크기에 따라 포인터를 이동시키면서 원하는 값을 찾을 수 있기 때문입니다.

val sumA = getSumArrayOf(pizzaA) // 연속합A
val sumB = getSumArrayOf(pizzaB) // 연속합B
sumA.sort() // 투 포인터를 사용하기 위해 정렬
sumB.sortDescending()

// 투포인터를 사용해서 정답을 찾는다
...

투포인터를 사용할 때 주의할 점이 있는데요. sumA + sumB = Target이 되는 두 포인터를 찾았을 때, 같은 연속합이 배열에 여러개 있을 수 있습니다. 따라서, 그때마다 중복된 값을 다 찾아서 곱한 값을 더해줘야 합니다.

while (indexA < arrayA.size && indexB < arrayB.size) {
        val sum = arrayA[indexA] + arrayB[indexB]

        if (sum < target) {
            indexA++
        } else if (sum > target) {
            indexB++
        } else { // 양쪽에서 중복 값을 다 찾아서 개수를 곱한다
            var countA = 0L
            var countB = 0L
            val prevA = arrayA[indexA]
            val prevB = arrayB[indexB]

            while (indexA < arrayA.size && arrayA[indexA] == prevA) {
                countA++
                indexA++
            }

            while (indexB < arrayB.size && arrayB[indexB] == prevB) {
                countB++
                indexB++
            }

            count += countA * countB
        }
    }

 

이분 탐색(upper_bound, lower_bound)

정렬된 데이터셋에서 원하는 값을 찾는 또 한가지 방법으로 이분 탐색이 있습니다. 여기서는 같은 연속합이 여러개 있을 수 있기 때문에 한쪽에서 연속합을 선택한 후, 반대편 연속합에서 원하는 값의 upper_bound와 lower_bound를 구해서 두 값을 빼주는 방식으로 target을 만드는 경우의 수를 구할 수 있습니다.

 

생각 과정

1. 각 피자에서 1조각 이상을 선택해서 특정 크기의 피자를 판매하려고 할 때, 가능한 경우의 수를 구하는 문제이다.
2. 피자조각을 선택할 땐 항상 연속해서 잘라야 하고, 한 쪽 피자만 사용할 수도 있다.
3. 우선 두 피자의 모든 연속합 필요하다.
4. 연속합을 모두 비교하면 시간 초과가 난다.
5. 구하려고 하는 합이 정해져있다 -> 각각의 연속합을 구해서 반대편에서 원하는 값이 있는지 확인하는 방식을 떠올린다. 각각의 연속합을 구하는 과정은 시간 내에 해결할 수 있다. 

 

코드1 (해싱)

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

fun main() = with(BufferedReader(InputStreamReader(System.`in`))) {
    val target = readLine().toInt()
    val (m, n) = readLine().split(" ").map(String::toInt)
    val pizzaA = IntArray(m) { readLine().toInt() }
    val pizzaB = IntArray(n) { readLine().toInt() }

    val sequentSumArray = getSumArrayOf(pizzaA)
    print(getCount(sequentSumArray, pizzaB, target))
}

// 1연속 조각 ~ n연속 조각의 합 개수 count (해싱)
fun getSumArrayOf(pizza: IntArray): IntArray {
    val sumCountArray = IntArray(1_000_001)
    var totalSum = 0 // 전체 합을 담을 변수

    for (i in 0 until pizza.size) {
        totalSum += pizza[i]
        var sum = 0

        for (j in i until i + pizza.size - 1) { // 전체 합은 제외
            sum += pizza[j % pizza.size] // 누적합, 원형 배열
            sumCountArray[sum]++
        }
    }

    sumCountArray[0] = 1 // 길이 0
    sumCountArray[totalSum] = 1 // 전체 합
    return sumCountArray
}

// pizzaB의 연속합을 구해서 pizzaA의 연속합에 target을 만드는 값이 있는지 확인
fun getCount(sequentSumArray: IntArray, pizzaB: IntArray, target: Int): Long {
    var count = 0L
    var totalSum = 0

    for (i in 0 until pizzaB.size) {
        totalSum += pizzaB[i]
        var sum = 0

        for (j in i until i + pizzaB.size - 1) {
            sum += pizzaB[j % pizzaB.size]

            // target을 만드는 sumA가 있는지 확인
            if (target - sum >= 0){
                count += sequentSumArray[target - sum]
            }
        }
    }

    count += sequentSumArray[target] // pizzaA만 사용하는 경우
    if (target - totalSum >= 0){     // pizzaB 전체 합
        count += sequentSumArray[target - totalSum]
    } 
    return count
}

 

코드2 (투 포인터)

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

fun main() = with(BufferedReader(InputStreamReader(System.`in`))) {
    val target = readLine().toInt()
    val (m, n) = readLine().split(" ").map(String::toInt)
    val pizzaA = IntArray(m) { readLine().toInt() }
    val pizzaB = IntArray(n) { readLine().toInt() }

    val sumA = getSumArrayOf(pizzaA)
    val sumB = getSumArrayOf(pizzaB)
    sumA.sort()
    sumB.sortDescending()

    print(getCount(sumA, sumB, target))
}

// 1연속 조각 ~ n연속 조각의 합을 담은 배열 반환
fun getSumArrayOf(pizza: IntArray): IntArray {
    val size = pizza.size
    val sumArray = IntArray(1 + size * (size-1) + 1)
    var sumIndex = 0
    var totalSum = 0 // 전체 합을 담을 변수

    // 연속합을 모두 저장
    for (i in 0 until pizza.size) {
        totalSum += pizza[i]
        var sum = 0

        for (j in i until i + pizza.size - 1) { // 전체 합은 제외
            sum += pizza[j % pizza.size]
            sumArray[sumIndex++] = sum // 누적합, 원형 배열
        }
    }

    sumArray[sumIndex++] = 0 // 길이 0
    sumArray[sumIndex] = totalSum // 전체 합
    return sumArray
}

// 투포인터
fun getCount(arrayA: IntArray, arrayB: IntArray, target: Int): Long {
    var count = 0L
    var indexA = 0
    var indexB = 0

    while (indexA < arrayA.size && indexB < arrayB.size) {
        val sum = arrayA[indexA] + arrayB[indexB]

        if (sum < target) {
            indexA++
        } else if (sum > target) {
            indexB++
        } else {
            var countA = 0L
            var countB = 0L
            val prevA = arrayA[indexA]
            val prevB = arrayB[indexB]

            while (indexA < arrayA.size && arrayA[indexA] == prevA) {
                countA++
                indexA++
            }

            while (indexB < arrayB.size && arrayB[indexB] == prevB) {
                countB++
                indexB++
            }

            count += countA * countB
        }
    }

    return count
}