[PS] 누적합

Hood·2024년 10월 16일

prifix sum ps

PS

목록 보기

2/15

✍ Kotlin으로 PS 문제 풀기

소속 중인 A&I 동아리에서 코딩 역량을 강화하고자
코딩캠프를 진행하며 작성한 포스트입니다.
해당 포스트는 Kotlin을 기반으로 작성하였습니다.

쿼리의 개념

연속된 숫자의 구간은 아래와 같이 표현할 수 있습니다.

1 ~ 10
2 ~ 30
66 ~ 100

예를 들어 1부터 N까지라면 다음과 같이 표현할 수 있습니다.

1 ~ N

또한 n부터 m까지라면 아래처럼 나타낼 수 있습니다.

n ~ m

이처럼 특정 구간을 나타내는 범위를 하나의 쿼리(Query)라고 정의할 수 있습니다.
PS 문제에서는 이러한 쿼리가 테스트 케이스마다 여러 개 주어지는 경우가 많습니다.
쿼리의 개수는 보통 문제의 입력 조건에서 확인할 수 있습니다.

부분합이란?

n부터 m까지의 쿼리가 주어졌을 때,
해당 구간의 모든 수의 합은 반복문으로 구할 수 있습니다.

예를 들어 n부터 m까지 직접 순회하며 더하면 되므로,
한 번의 쿼리를 처리하는 데 걸리는 시간 복잡도는 구간 길이에 비례합니다.

즉, 시간 복잡도는 대략 O(M) 또는 O(m - n + 1) 로 볼 수 있습니다.

그렇다면 쿼리가 여러 개 주어진다면 어떨까요?

예를 들어 n부터 m까지의 합을 구하는 쿼리가 k번 등장한다면,
매번 반복문을 새로 돌려야 하므로 전체 시간 복잡도는
O(K × M) 이 됩니다.

쿼리의 개수가 많아질수록 시간 복잡도는 빠르게 커집니다.
수의 범위가 1,000,000이고, 이를 100,000번 반복해서 계산해야 한다면
이미 시간 초과가 발생할 가능성이 매우 높습니다.

누적합 알고리즘

이럴 때 사용하는 대표적인 방법이 바로 누적합(Prefix Sum) 알고리즘입니다.
쉽게 말해, 미리 합을 계산해 둔 배열을 만들어 두는 방법입니다.

예를 들어 1부터 9까지의 배열이 있다고 가정해 보겠습니다.
누적합 배열은 1부터 현재 위치까지의 합을 순서대로 저장합니다.

일반적으로 누적합 배열은 앞에 0을 하나 두고 시작합니다.
이렇게 하는 이유는 구간 합을 계산할 때 인덱스 처리를 더 단순하게 만들기 위해서입니다.

예를 들어 prefixSum[i]를 1부터 i까지의 합이라고 하면,
l부터 r까지의 구간 합은 다음과 같이 계산할 수 있습니다.

prefixSum[r] - prefixSum[l - 1]

앞에 0이 없다면 l = 1인 경우를 따로 처리해야 해서 코드가 복잡해질 수 있습니다.

예를 들어 4부터 9까지의 합을 구하고 싶다면,
이미 구해 둔 누적합에서 다음처럼 계산하면 됩니다.

prefixSum[9] - prefixSum[3]

즉, 끝 위치까지의 합에서 시작 직전까지의 합을 빼는 방식입니다.

이렇게 하면 처음 누적합 배열을 만드는 데 O(M),
이후 각 쿼리를 처리하는 데 O(1) 이 걸립니다.

따라서 쿼리가 K개라면 전체 시간 복잡도는

O(M + K)

가 됩니다.

예를 들어 수의 범위가 1,000,000이고 쿼리가 100,000개라면,
기존의 방식은 거의

O(1,000,000 \times 100,000)

수준이지만, 누적합을 사용하면

O(1,000,000 + 100,000) = O(1,100,000)

으로 크게 줄어듭니다.

2차원 누적합 알고리즘

누적합은 1차원뿐만 아니라 2차원 배열에서도 사용할 수 있습니다.
다만 2차원에서는 계산 과정이 조금 더 복잡해집니다.

예를 들어 5 × 5 크기의 2차원 배열이 있다고 가정해 보겠습니다.
이때 (1, 1)부터 (2, 2)까지의 직사각형 영역의 합을 구하고 싶다면,
단순히 위쪽과 왼쪽 누적합만 더하면 겹치는 부분이 생깁니다.

그래서 다음과 같은 방식으로 계산합니다.

pSum[x][y] = pSum[x-1][y] + pSum[x][y-1] - pSum[x-1][y-1] + arr[x][y]

즉,

위쪽 누적합을 더하고
왼쪽 누적합을 더한 뒤
두 번 더해진 왼쪽 위 영역을 한 번 빼고
현재 값을 더하는 방식입니다

이유는 pSum[x-1][y]와 pSum[x][y-1]를 더할 때
왼쪽 위 영역인 pSum[x-1][y-1]이 중복으로 포함되기 때문입니다.
따라서 한 번 빼 주어야 올바른 값이 됩니다.

그렇다면 (2, 2)부터 (4, 4)까지의 합은 어떻게 구할까요?

이 경우에는 다음 공식을 사용합니다.

pSum[4][4] - pSum[1][4] - pSum[4][1] + pSum[1][1]

일반화하면 (x1, y1)부터 (x2, y2)까지의 구간 합은 다음과 같습니다.

pSum[x2][y2] - pSum[x1-1][y2] - pSum[x2][y1-1] + pSum[x1-1][y1-1]

이 식을 사용하면 직사각형 영역의 합도 O(1) 에 구할 수 있습니다.

in Kotlin

아래 코드는 2차원 누적합 배열을 만드는 예시입니다.

import java.io.StreamTokenizer

fun main() = with(StreamTokenizer(System.`in`.bufferedReader())) {
    fun nextInt(): Int {
        nextToken()
        return nval.toInt()
    }

    val x = nextInt()
    val y = nextInt()

    val pSum = Array(x + 1) { IntArray(y + 1) }

    repeat(x) { i ->
        repeat(y) { j ->
            pSum[i + 1][j + 1] = pSum[i + 1][j] + pSum[i][j + 1] - pSum[i][j] + nextInt()
        }
    }

    pSum.forEach { println(it.toList()) }
}

코드에서는 x와 y를 입력받아
(1, 1)부터 각 위치까지의 누적합을 pSum 배열에 저장합니다.
이렇게 미리 누적합 배열을 만들어 두면, 이후 여러 구간의 합을 매우 빠르게 구할 수 있습니다.

만약 (2, 2)부터 (4, 4)까지의 값을 구하고 싶다면
다음과 같은 코드를 사용할 수 있습니다.

val sb = StringBuilder()

repeat(q) {
    val x1 = nextInt()
    val y1 = nextInt()
    val x2 = nextInt()
    val y2 = nextInt()

    sb.appendLine( pSum[x2][y2] - pSum[x2][y1 - 1] - pSum[x1 - 1][y2] + pSum[x1 - 1][y1 - 1]
    )
}

StringBuilder를 사용하면 출력이 많은 경우에도 더 효율적으로 처리할 수 있습니다.

그리고 핵심은 아래 식입니다.

pSum[x2][y2] - pSum[x2][y1 - 1] - pSum[x1 - 1][y2] + pSum[x1 - 1][y1 - 1]

이 식을 사용하면 (x1, y1)부터 (x2, y2)까지의 구간 합을
반복문 없이 빠르게 구할 수 있습니다.

📌 결론

상위 PS 문제를 풀다 보면 누적합은 기본적으로 이해하고 있어야 하는 개념입니다.
특히 구간 합을 여러 번 물어보는 문제에서는 거의 필수적으로 등장합니다.

정리해 보면 다음과 같습니다.

단순 반복으로 구간 합을 구하면 쿼리 수가 많을 때 비효율적입니다.
누적합을 사용하면 구간 합을 훨씬 빠르게 구할 수 있습니다.
2차원 누적합까지 이해하면 표, 격자, 보드 형태의 문제에도 응용할 수 있습니다.

누적합은 처음 보면 다소 헷갈릴 수 있지만,
공식을 직접 손으로 써 보고 예제를 따라가 보면 금방 익숙해집니다.
PS에서 자주 등장하는 만큼, 꼭 익혀 두시면 큰 도움이 됩니다.

Hood

달을 향해 쏴라, 빗나가도 별이 될 테니 👊

이전 포스트

[PS] 소수 판별 알고리즘

다음 포스트