구간 합(Prefix Sum) + 문제

정순동·2024년 3월 7일

알고리즘

목록 보기
32/33

구간 합이란?

프로그래밍을 하다 보면 '구간 합'과 '부분 합'에 대해 들어볼 수 있는데 이의 뜻은 각각 아래와 같다.

  • 구간 합(Prefix Sum) : 구간 합이란 수들의 나열에서 특정 구간의 합을 의미한다. 구간 합 알고리즘은 보통 1차원 배열에서 i ~ k 인덱스 사이의 값들의 합을 구하는데 사용한다.

  • 부분 합(Partial Sum) : 부분 합은 구간 합과는 달리 처음부터 특정 인덱스까지의 합을 의미한다. 즉 인덱스가 0 ~ k 인 요소의 합을 말한다.

포함 관계에 있다고 보면 된다.


구간 합의 핵심이론

구간 합은 합 배열을 이용해 시간 복잡도를 더 줄이기 위해 사용하는 특수 목적 알고리즘이다. 사용 빈도가 높으니 잘 알아두는게 좋다.

합배열이란?

구간 합 알고리즘을 활용하러면 먼저 합 배열을 구해야 한다. 합 배열에 대한 설명은 글로 읽는것 보다는 아래 정의 방법과 그림을 통해 알아보는 것이 훨씬 이해하기 편하다.

↓ 합 배열 S 정의 방법

	int[] A = {15,13,10,7,3,12};
    int[] S = new int[A.length]; // A와 같은 크기인 배열 S생성
    S[0] = A[0]; // 첫 요소만 따로 지정
	for (int i = 1; i = A.length; i++) { // 첫 요소 제외 1인덱스부터 시작
    	S[i] = S[i-1] + A[i];
        // 예를 들어 S[2] = A[0] + A[1] + A[2]가 되는것임.
    }

합 배열은 기존의 배열을 전처리한 배열이라 생각하자.

아래 그림을 통해 무슨말인지 알아보자.

이렇게 합 배열이란 배열 A의 모든 인덱스에서의 부분 합을 구해서 배열 S의 각 인덱스 에 맞게 저장해 둔 배열이다.
S[0] = A[0]
S[1] = S[0] + A[1]
S[2] = S[1] + A[2]
S[3] = S[2] + A[3]
...
이런 식으로 말이다.

여기서 S[0] = A[0] 처럼 따로 지정하기 싫은 사람들은 A,S배열의 크기를 1칸씩 늘려 아예 0번째 인덱스는 사용하지 않는 방법도 있다.

이렇게 합 배열을 만들어 두면 후에 배열의 어느 구간의 합을 구하려고 계속 더했다가 뺏다가 할 필요없이 합 배열의 해당 부분만 가져오면 된다.

따라서 합 배열을 만드는 과정 자체는 복잡도가 O(n)이다. 하지만 합 배열은 단 한번만 만들면 되기에 그로인한 이점이 더 크다.

구간 합을 구하는 공식

합 배열을 미리 구해놓으면 기존 배열의 일정 범위의 합(구간 합)을 구하는 시간 복잡도가 O(N)에서 O(1)로 감소하게 된다.

아래 예시를 통해 알아보자.

A 배열 2~5요소의 합을 구하는 방법이 무엇일까? 2가지로 나눠 볼 수 있겠다.
1. A[2] + A[3] + A[4] + A[5] = 32
2. S[5] - S[2] = 32

2번 즉, 합 배열을 활용해서 구간 합을 구하는 방식이 훨씬 간단하다.
요소의 크기가 대단히 커졌을 때에도 O(1), 무조건 한방에 끝나는 구간 합 공식이다.


합 배열과 구간 합 공식을 적재적소에 활용하면 코딩 테스트나 자료구조 최적화 등에서 큰 이점으로 작용할 것이다.

문제

구간 합 구하기 - 백준11659번

수 N개가 주어졌을 때 i번째 수에서 j 번째 수까지의 합을 구하는 프로그램 작성

입력│←
1번째 줄에 수의 개수 N(1 <= N <= 100,000), 합을 구해야 하는 횟수 M(1 <= M <= 100,000), 2번째 줄에 N개의 수가 주어진다. 각 수는 1,000보다 작거나 같은 자연수이다. 3번째 줄부터는 M개의 줄에 합을 구해야 하는 구간 i와 j가 주어진다.

출력│→
총 M개의 줄에 입력으로 주어진 i번째 수에서 j번째 수까지의 합을 출력한다.

추가 사항
제한 시간 : 0.5초
난이도 : 실버3


문제 분석하기

문제에서 수의 개수, 합을 구해야 하는 횟수는 최대 100,000개/번이다.
따라서 만약 100,000개의 수와 100,000번의 합을 구할 때 10만번의 구간 합을 모두 따로 계산한다면 제한 시간인 0.5초 이내에 절대 풀 수 없는 문제이다.

※ 최악의 경우 1억번을 훨씬 넘는 연산을 수행하게 되어 0.5초이내 풀 수 없다. 0.5초가 기준이라면 최대 5천만번 이내의 연산을 수행해야 한다.

따라서 합 배열을 하나 만들어서 합 배열을 통해 구간 합을 구하도록 하자.

코드

public class 구간합11659 {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        int suNo = Integer.parseInt(st.nextToken());
        int sumNo = Integer.parseInt(st.nextToken());

        long[] S = new long[suNo + 1];
        st = new StringTokenizer(br.readLine());
        
        for (int i = 1; i < S.length; i++) {
            S[i] = S[i - 1] + Integer.parseInt(st.nextToken());
        }
        
        for (int s = 0; s < sumNo; s++) {
            st = new StringTokenizer(br.readLine());
            int i = Integer.parseInt(st.nextToken());
            int j = Integer.parseInt(st.nextToken());
            System.out.println(S[j] - S[i-1]);
        }
    }
}

위의 설명한 내용만 알고있으면 코드는 간단하게 풀 수 있을 것이다.


구간 합 구하기 2 - 백준11660번

N x N개의 수가 N x N 크기의 표에 채워져 있다. 표 안의 수 중 (X₁,Y₁)에서 (X₂,Y₂)까지의 합을 구하려 한다. X는 행이고 Y는 열을 의미한다.

예를 들어 N = 4이고, 표가 다음과 같이 채워져 있을 때를 살펴보자.

(2,2)에서 (3,4)까지의 합을 구하면 3+4+5+4+5+6 = 27이고, (4,4)에서 (4,4)까지의 합을 구하면 7이다. 표에 채워져 있는 수와 합을 구하는 연산이 주어졌을 때 이를 처리하는 프로그램을 작성하시오.

입력│←
첫째 줄에 표의 크기 N과 합을 구해야 하는 횟수 M이 주어진다. (1 ≤ N ≤ 1024, 1 ≤ M ≤ 100,000)
둘째 줄부터 N개의 줄에는 표에 채워져 있는 수가 1행부터 차례대로 주어진다. 다음 M개의 줄에는 네 개의 정수 x1, y1, x2, y2 가 주어지며, (x1, y1)부터 (x2, y2)의 합을 구해 출력해야 한다. 표에 채워져 있는 수는 1,000보다 작거나 같은 자연수이다. (x1 ≤ x2, y1 ≤ y2)

출력│→
총 M줄에 걸쳐 (x1, y1)부터 (x2, y2)까지 합을 구해 출력한다.

추가 사항
제한 시간 : 1초
난이도 : 실버1


문제 분석하기

먼저 질의의 개수가 100,000이므로 이 문제 역시 질의마다 합을 구하면 절대 안된다.

따라서 구간 합 배열을 이용해야 한다는 것을 알 수 있는데 앞선 문제와 다른 점이라면, 배열이 1차원 → 2차원으로 확장 됐다는 점이다.
2차원 배열으로 어떻게 구간 합 배열을 구성할지 고민하는 것이 이 문제의 핵심이다.

2차원 구간 합 배열 D[x][y]의 정의

D[x][y] = 원본 배열 (0, 0)부터 (x, y)까지의 사각형 영역 안에 있는 수의 합

위 표를 예시로 보면 D[2][2] = 8이 돼야 한다.

따라서 D[i][j]를 채우는 공식은아래와 같다.
D[i][j] = D[i][j-1] + D[i-1][j] - D[i-1][j-1] + A[i][j]

왜? 아래 그림을 보면서 D[3][3] 부분을 구한다고 해 보자.

왼쪽 합 배열D의 빨간색 15는 오른쪽 기존 배열의 빨간부분의 합을 의미하는 수이다.
왼쪽 합 배열D의 주황색 15는 오른쪽 기존 배열의 주황부분의 합을 의미하는 수이다.
그리고 왼쪽 합 배열D의 노란색 8은 오른쪽 기존 배열의 노란부분의 합을 의미하는 수이다.
또한 노란부분은 빨간부분과 주황부분이 서로 겹치게 되는 부분이기도 하다.
보라색부분은 그저 A[3][3]의 값인 5만을 의미한다

우리가 구하고자 하는 D[3][3]은 오른쪽 기존배열에서 초록색부분의 합을 의미하는데 그렇다면 이는
주황색부분 + 빨강색부분 - 노란색부분(겹치는부분) + 보라색부분이 된다.

따라서 우리는 공식이 'D[i][j] = D[i][j-1] + D[i-1][j] - D[i-1][j-1] + A[i][j]' 임을 알 수 있다.
D[3][3] = D[3][2] + D[2][3] - D[2][2] + A[3][3]으로 D[3][3]은 27이된다

이런식으로 2차원 합 배열을 완성할 수 있다.

합 배열을 통해 질의에 답하는 방법

합 배열을 통해 2 2 3 4 라는 질의에 답 해보자. (2,2) (3,4)까지의 합을 구해야한다.

해당 부분을 어떻게 구할까? 답은 위에서 합 배열을 만들 때 처럼 하면된다.
아래 그림을 보면서 이해해보자.

우리가 구하고자 하는 부분은 초록색부분이다.
따라서 빨간색부분 - 주황색부분 - 파란색부분 + 보라색부분(겹치는부분) 이다.
위 합 배열을 만드는 방법과 매우 비슷하다. 보라색부분이 겹쳐 2번빠지는것을 고려해서 마지막에 추가해 주면 끝이다.

이를 합 배열을 통해 선택하자면 아래 그림처럼 된다

똑같이 빨강 - 주황 - 파랑 + 보라 해 주면 우리가 원하는 값인 27을 얻을 수 있다.

질의 X₁,Y₁,X₂,Y₂를 구하는 방법은
'D[X₂][Y₂] - D[X₁-1][Y₂] - D[X₂][Y₁-1] + D[X₁-1][Y₁-1]'
이 된다.

코드

public class 구간합11660 {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int N = Integer.parseInt(st.nextToken());
        int M = Integer.parseInt(st.nextToken());

        int A[][] = new int[N + 1][N + 1];
        for (int i = 1; i <= N; i++) {
            st = new StringTokenizer(br.readLine());
            for (int j = 1; j <= N; j++) {
                A[i][j] = Integer.parseInt(st.nextToken());
            }
        }

        int D[][] = new int[N + 1][N + 1];
        for (int i = 1; i <= N; i++) {
            for (int j = 1; j <= N; j++) {
                // 구간 합 구하는 공식
                D[i][j] = D[i-1][j] + D[i][j-1] - D[i-1][j-1] + A[i][j];
            }
        }

        // 질의 응답 코드
        for (int i = 0; i < M; i++) {
            st = new StringTokenizer(br.readLine());
            int x1 = Integer.parseInt(st.nextToken());
            int y1 = Integer.parseInt(st.nextToken());
            int x2 = Integer.parseInt(st.nextToken());
            int y2 = Integer.parseInt(st.nextToken());
            // 질의에 답변하는 코드
            int result = D[x2][y2] - D[x1-1][y2] - D[x2][y1-1] + D[x1-1][y1-1];
            System.out.println(result);
        }
    }
}

어려운 부분은 없을것이다.


나머지 합 구하기 - 백준10986번

N개의 수 A₁,A₂...,An이 주어졌을 때 연속된 부분의 합이 M으로 나누어떨어지는 구간의 개수를 구하는 프로그램을 작성하시오. 즉, Ai + ... + Aj(i <= j)의 합이 M으로 나누어 떨어지는 (i,j) 쌍의 개수를 구하시오

입력│←
첫째 줄에 N과 M이 주어진다. (1 ≤ N ≤ 10⁶, 2 ≤ M ≤ 10³)

둘째 줄에 N개의 수A₁,A₂...,An이 주어진다. (0 ≤ Ai ≤ 10⁹)

출력│→
1번째 줄에 연속된 부분의 합이 M으로 나누어 떨어지는 구간의 개수를 출력한다.

추가 사항
제한 시간 : 1초
난이도 : 골드3
성공률 : 26%

문제 분석하기

N이 최대 10⁶, M이 최대 10³이므로 1억번 이내의 연산 횟수로 끝내야 한다 따라서 위 2문제들과 마찬가지로 합 배열을 사용해서 문제를 풀어야 한다.

나머지 합 문제 풀이의 핵심 아이디어

  • (A + B) % C 는 ((A % C) + (B % C)) % C와 같다. 다시 말해 특정 구간 수들의 나머지 연산을 더해 나머지 연산을 한 값과 이 구간 합의 나머지 연산을 한 값은 동일하다.
  • 구간 합 배열을 이용한 식 S[j] - S[i]는 원본 배열의 i + 1부터 j까지의 구간 합이다.
  • S[j] % M의 값과 S[i] % M의 값이 같다면 (S[j] - S[i]) % M은 0이다. 즉, 구간 합 배열의 원소를 M으로 나눈 나머지로 업데이트하고 S[j]와 S[i]가 같은 (i,j)쌍을 찾으면 원본 배열에서 i + 1부터 j까지의 구간 합이 M으로 나누어떨어짐을 알 수 있다.

이해가 안된다면 아래 예시들을 통해 알아보면 더 쉬울 것이다.

새로운 합 배열 만들기

원본 배열 A를 통해 합 배열 S를 만들고, 해당 합 배열 S의 각 요소를 M으로 나눈 나머지를 저장한다.

오른쪽 아래 변경된 합배열 S를 가지고 문제를 풀 것이다.

  1. 우선 변경된 합 배열S에서 원소 값이 0인 개수만 세어 정답에 미리 더해둔다. 변경된 합 배열의 원소 값이 0이라는 뜻은 원본 배열에 0부터 i까지의 구간 합이 이미 M으로 나누어 떨어지는 값이기 때문이다.

  2. 이제 변경된 합 배열S에서 원소 값이 같은 인덱스의 개수, 즉 나머지 값이 같은 합 배열의 개수를 센다. 변경된 합 배열에서 원소값이 같은 2개의 원소를 뽑는 모든 경우의 수를 구해 정답에 더하면 된다. 위의 예시에서는 0이 3개, 1이 2개이므로 ₃C₂, ₂C₂로 경우의 수를 구하여 더한다.

₃C₂ → 총 경우의 수 +3
₂C₂ → 총 경우의 수 +1
따라서 총 경우의 수 = 3(1번) + 3 + 1 = 7

부가설명 ↓

조합의 경우에 변경된 합 배열 S에서 요소 '1'를 선택해서 경우의 수를 알아보자.

일단 요소'1'은 변경된 합 배열 S에서 두번 이상 나오니 조합이 가능하다.
같은 이유로 요소'0'도 변경된 합 배열 S에서 3번 나오니 조합이 가능하다.

S[0] = 1이고 S[3] = 1이라는것의 의미는 S[3] - S[0]을 할 시에 0이 된다는 것인데, 그 말은 원본배열에서 각각 해당하는 부분을 서로 빼고 남는 부분의 % M = 0이 된다는 의미이다.

오른쪽 그림의 빨간부분 - 주황부분 = 초록부분, 초록부분의 합 % M = 0이다.

A[0~4] = 7, A[0] = 1
7 - 1 = 6
6 % M(3) = 0

그래서 위 경우의 수가 정답에 카운트되는 것이다.

변경된 배열S에서 0은 3개가있다 따라서 3개중 2개를 골라 서로 빼 주고
남은 부분을 % M 했을 때 나머지는 무조건 0이된다.

코드

public class 나머지합10986 {
    public static void main(String[] args) throws IOException {
        Scanner sc = new Scanner(System.in);
        int N = sc.nextInt();
        int M = sc.nextInt();
        long[] S = new long[N];
        long[] C = new long[N];
        long answer = 0;
        S[0] = sc.nextInt();

        for (int i = 1; i < N; i++) { // 수열 합 배열 만들기
            S[i] = S[i-1] + sc.nextInt();
        }

        for (int i = 0; i < N; i++) { // 합 배열의 모든 값에 %연산을 수행사여 새로운 합 배열 S 만들기
            int remainder = (int) (S[i] % M);
            // 0 ~ i까지의 구간 합 % M이 0일 때 정답에 더하기
            if (remainder == 0) answer++;
            // 나머지가 같은 인덱스의 개수 카운팅하기
            C[remainder]++;
        } // C[3,2,0,0,0] 나머지 값 0이 3개, 나머지 값 1이 2개라는 뜻.

        for (int i = 0; i < M; i++) {
            if (C[i] > 1) { // 나머지가 같은 인덱스가 2개 이상인가?
                // 나머지가 같은 인덱스 중 2개를 뽑는 경우의 수를 더하기
                answer = answer + (C[i] * (C[i] - 1) / 2); 
                // N개 중 2개를 뽑는 조합 공식임
            }
        }
        System.out.println(answer);
    }
}

문제의 예제를 기준으로 배열 C에는 [3,2,0,0,0]이 들어가게 된다.
C[0] = 3이라는 의미는 나머지 값이 0인 요소가 3개라는 뜻이고,
C[1] = 2이 가지는 의미는 나머지 값이 1인 요소가 2개라는 뜻이다.

이를 바탕으로 맨 아래 for문에서 나머지가 같은 요소가 2개 이상인지를 확인하고 2개 이상이라면 N개중 2개를 뽑는 조합 공식을 사용해서 answer에 답을 누적한다.

문제의 예제를 기준으로 맨 아래 for문에서 처음 C[0]은 3이다.
if(3 > 1)이 통과되고 ₃C₂ 를 계산 해 준다. 3개의 요소중 2개를 뽑아 만들 수 있는 조합의 개수라는 뜻이다. 3 x 2 / 2가 되고 answer에는 3이 추가된다.

for문이 2번째 돌 때 C[1]은 2이다. 나머지가 1인 요소가 2개라는 뜻이고
if(2 > 1)이 통과되며 마찬가지로 ₂C₂를 계산한다. 2 x 1 / 2가 되며 answer에 1을 추가한다.

나머지 C[2,3,4]는 값이 0이라 for문을 더 돌아도 if문을 통과하지 못한다. 값이 1이어도 의미 없다.

answer에 답이 추가됐던 경우는 아래와 같다

  • 배열S 의 각 요소 중 값이 0인 요소의 개수만큼 answer에 추가.
    S[1],S[2],S[4]의 각 요소가 0이었으니 answer에 3추가.
  • ₃C₂ 를 계산한 결과를 answer에 추가.
    계산 결과는 3이니 answer에 3추가.
  • ₂C₂ 를 계산한 결과를 answer에 추가.
    계산 결과는 1이니 answer에 1추가.

따라서 예제 기준 총 경우의 수는 7번이되겠다.

생각보다 설명하기 복잡해서 그렇지 한번 이해하면 성공률 26% 문제 치고는 어렵지 않다. 다만 앞서 말했든 제한 시간이 1초기 때문에 무조건 합 배열을 사용해야 한다는 점만 유의하면 된다.

0개의 댓글