자료구조 - 구간합

eunsong·2024년 11월 26일

algorithm

Algorithm

목록 보기

3/6

구간 합은 합 배열을 이용하여 특정 구간의 합을 구할 때 시간 복잡도를 더 줄이기 위해 사용하는 특수한 목적의 알고리즘이다. 코딩 테스트에서 유용하게 활용 가능하니 꼭 알고 넘어가야 하는 알고리즘 중 하나이다.

구간합 정의

구간합은 배을의 특정 구간 [i, j] 의 합을 빠르게 구하기 위한 기법이다.
ex) 배열 [2, 4, 5, 7, 8] 에서 구간 [1, 3]의 합은 4+5+7=16.

문제점?

단순히 반복문을 이용해 구간합을 구하면 O(n)의 시간 복잡도가 걸리기 때문에 여러 번 구간합을 구해야 하는 경우 비효율적이다.

해결책

누적합(Prefix Sum)을 사용해 구간합을 O(1)에 계산할 수 있다.

fun calculatePrefixSum(arr: IntArray): IntArray {
    val prefixSum = IntArray(arr.size)
    prefixSum[0] = arr[0]
    for (i in 1 until arr.size) {
        prefixSum[i] = prefixSum[i - 1] + arr[i]
    }
    return prefixSum
}

fun getRangeSum(prefixSum: IntArray, start: Int, end: Int): Int {
    return if (start == 0) prefixSum[end] else prefixSum[end] - prefixSum[start - 1]
}

fun main() {
    val arr = intArrayOf(2, 4, 5, 7, 8)
    val prefixSum = calculatePrefixSum(arr)
    println("Prefix Sum Array: ${prefixSum.joinToString()}")
    println("Sum of range [1, 3]: ${getRangeSum(prefixSum, 1, 3)}") // Output: 16
}

누적합을 미리 계산하는 과정은 O(N) 의 시간복잡도를 가지고 있지만 중요한 점은
한 번만 계산하면 이후의 구간합 계산은 O(1)이라는 것이다.

초기 누적합 계산: O(N)

배열의 크기 N에 비례해 누적합 배열을 만드는 데 걸리는 시간
sum[i] = arr[0] + arr[1] + ... + arr[i]

구간합 쿼리: O(1)

한번 계산된 누적합 배열을 사용해 [i, j] 구간의 합을 바로 계산할 수 있음
s(i, j) = sum[j] - sum[i - 1]

누적합의 한계

배열 값이 업데이트되면 전체를 다시 계산해야 한다.

세그먼트 트리(Segment Tree)

값을 업데이트하거나 구간합을 구할 때도 효율적으로 동작(O(logn))

2차원 구간합

2차원 배열에서의 누적합과 구간합을 구하는 방법은 1차원 배열의 구간합 알고리즘을 응용하여 구할 수 있다.

크기 N x M 의 2차원 배열에서 특정 영역 (x1, y1) ~ (x2, y2) 까지의 합을 여러 번 계산해야 할 떄, 매번 반복문을 돌리면 비효율적이다.

반복문을 사용하지 않고 O(1) 로 직사각형 영역의 합을 구합니다.

누적합 배열 prefixSum[i][j] 은 (1, 1) ~ (i, j) 까지의 모든 원소의 합을 나타냅니다.

prefixSum[i][j]=arr[1][1]+⋯+arr[i][j]

2차원 누적합 배열 만들기

arr[i][j]: 현재 위치의 값.
prefixSum[i-1][j]: 현재 행 바로 위의 누적합.
prefixSum[i][j-1]: 현재 열 바로 왼쪽의 누적합.
prefixSum[i-1][j-1]: 중복 계산된 값을 빼준다.

fun calculate2DPrefixSum(arr: Array<IntArray>): Array<IntArray> {
    val rows = arr.size
    val cols = arr[0].size
    val prefixSum = Array(rows) { IntArray(cols) }

    for (i in 0 until rows) {
        for (j in 0 until cols) {
            val current = arr[i][j]
            val top = if (i > 0) prefixSum[i - 1][j] else 0
            val left = if (j > 0) prefixSum[i][j - 1] else 0
            val topLeft = if (i > 0 && j > 0) prefixSum[i - 1][j - 1] else 0
            prefixSum[i][j] = current + top + left - topLeft
        }
    }
    return prefixSum
}

직사각형 영역의 합 구하기

누적합 배열을 이용해 영역 (x1, y1)부터 (x2, y2)까지의 합을 계산합니다.

sum = prefixSum[x2][y2] − prefixSum[x1−1][y2] − prefixSum[x2][y1−1] + prefixSum[x1−1][y1−1]

ex) 아래와 같이 (2, 2) ~ (3, 4) 일 경우,
(3, 4) 구간 합에서 (1, 4) 구간 합, (3, 1) 구간 합을 뺀 다음 중복하여 뺀 (1, 1) 을 더하면 됩니다.

prefixSum[3][4] - prefixSum[1][4] - prefixSum[3][1] + prefixSum[1][1] = 27

fun get2DRangeSum(
    prefixSum: Array<IntArray>, x1: Int, y1: Int, x2: Int, y2: Int
): Int {
    val total = prefixSum[x2][y2]
    val top = if (x1 > 0) prefixSum[x1 - 1][y2] else 0
    val left = if (y1 > 0) prefixSum[x2][y1 - 1] else 0
    val topLeft = if (x1 > 0 && y1 > 0) prefixSum[x1 - 1][y1 - 1] else 0
    return total - top - left + topLeft
}

fun main() {
    val arr = arrayOf(
        intArrayOf(1, 2, 3),
        intArrayOf(4, 5, 6),
        intArrayOf(7, 8, 9)
    )

    val prefixSum = calculate2DPrefixSum(arr)
    println("Prefix Sum Array:")
    prefixSum.forEach { println(it.joinToString()) }

    println("Sum of range [(1, 1) to (2, 2)]: ${get2DRangeSum(prefixSum, 1, 1, 2, 2)}") // Output: 28
}