11660 - 구간 합 구하기 5

이 문제는 2차원 누적 합을 잘 사용해야 하는 문제이다. 문제 내용을 살펴보면, $N \times N$ 크기의 2차원 표를 입력받아 $M$개의 2차원 구간합을 구하는 것이 목표이다. 그런데 만약 표를 입력받아 순회하며 합을 구하는 브루트포스 방식으로 단순히 구현하면 최악의 경우 100000 * 1024 * 1024번(1천억 번) 연산하므로 시간초과가 날 것이다.

따라서 앞서 말했듯 이 문제를 풀려면 2차원 누적 합을 사용해야 한다. 누적합 테이블의 $(i, j)$ 번째 원소를 $(1,1)$ 부터 $(i, j)$ 까지 모든 원소들의 합으로 정의하자. 그렇다면 $(i, j)$ 번째 누적합 원소는 다음과 같이 얻을 수 있다.

여기서 마지막 항에 $- \text{prefixSum}[i-1][j-1]$가 붙은 이유는 $\text{prefixSum}[i-1][j] + \text{prefixSum}[i][j-1]$ 을 더하면 prefixSum의 $(i-1, j-1)$번째 원소는 두 번 더해지기 때문이다. 이는 마치 집합에서 합집합의 원소 개수를 구할 때 쓰는 포함-배제 원리와 유사하다고 생각하면 된다.

그림으로 보면 위와 같다.

그렇다면 문제에서 원하는, 특정 구간의 합을 얻을 때는 prefixSum 테이블을 어떻게 사용해야 할까? 먼저 prefixSum에서 (i, j)번째 원소는 (1,1)에서 (i,j) 까지의 구간합이다. 문제는 (x1, y1) ~ (x2, y2)의 원소들의 합을 원한다. prefixSum[x2][y2]는 (1,1) ~ (x2, y2)의 값을 줄 것이다. 해당 사각형에서 (x1-1, y2)와 (x2, y1-1)을 제외하면 우리가 원하는 구간합에서 prefixSum[x1-1][y1-1] 만큼 적은 값을 얻는다. prefixSum을 얻을 때처럼 그 부분이 겹치기 때문이다. 따라서 (x1-1, y1-1)를 한 번 더해줘야 한다. 이를 수식으로 적으면 다음과 같다

🔗 문제 링크

코드

미리 구해둔 식을 바탕으로 코드를 작성하면 된다. 먼저 필요한 변수를 정의한다. table는 입력받은 테이블의 원소를 저장해둘 배열이며. pSum은 누적합을 저장해둘 배열이다.

class Main {
    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
    static int N, M;
    static int[][] table;
    static int[][] pSum;
  // ...

다음으로 main 함수에서 N, M을 입력받고 table과 pSum을 메모리에 할당한다. 배열 크기를 N+5로 둔 이유는 먼저 1 based indexing을 하기 위함과, 혹시나 모를 실수로 배열에 잘못 접근해 OutOfBound 예외가 나는 것을 방지하기 위함이다. 참고로 1 based indexing을 하는 이유는 앞선 수식들에서 i-1, j-1을 계산할 때 i나 j가 0일 경우 OOB가 날 수 있어 예외처리를 따로 해줘야 하는 불편함이 있기 때문이다. 하지만 1 based indexing을 하면 i나 j가 1이어도 OOB가 나지 않고, arr[i][0], arr[0][j]의 값은 0으로 자동 초기화되어 있기 때문에 예외처리를 할 필요가 없다.

public static void main(String[] args) throws IOException {
    StringTokenizer st = new StringTokenizer(br.readLine());
    N = Integer.parseInt(st.nextToken());
    M = Integer.parseInt(st.nextToken());

    table = new int[N+5][N+5];
    pSum = new int[N+5][N+5];
  // ...

이제 테이블의 원소를 입력받아 넣어둔 뒤 누적합을 계산한다. 첫 번째 루프에서 원소를 입력받은 뒤 두 번째 루프에서 누적합을 계산한다. 루프 안의 계산식은 앞서 살펴보았던 누적합 계산식이다.

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

for (int i = 1; i <= N; i++) {
    for (int j = 1; j <= N; j++) {
        pSum[i][j] = table[i][j] + pSum[i-1][j] + pSum[i][j-1] - pSum[i-1][j-1];
    }
}

마지막으로 M번 반복하면서 구간을 입력받아 구간합 계산식을 적용한 뒤 출력한다.

    while (M-- > 0) {
        st = new StringTokenizer(br.readLine());
        int x1, y1, x2, y2;
        x1 = Integer.parseInt(st.nextToken());
        y1 = Integer.parseInt(st.nextToken());
        x2 = Integer.parseInt(st.nextToken());
        y2 = Integer.parseInt(st.nextToken());

        int result = pSum[x2][y2] - pSum[x1-1][y2] - pSum[x2][y1-1] + pSum[x1-1][y1-1];
        bw.write(result + "\n");
    }

    bw.flush();
    bw.close();
}

전체 코드

import java.util.*;
import java.io.*;

class Main {
    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
    static int N, M;
    static int[][] table;
    static int[][] pSum;

    public static void main(String[] args) throws IOException {
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        M = Integer.parseInt(st.nextToken());

        table = new int[N+5][N+5];
        pSum = new int[N+5][N+5];


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

        for (int i = 1; i <= N; i++) {
            for (int j = 1; j <= N; j++) {
                pSum[i][j] = table[i][j] + pSum[i-1][j] + pSum[i][j-1] - pSum[i-1][j-1];
            }
        }
        while (M-- > 0) {
            st = new StringTokenizer(br.readLine());
            int x1, y1, x2, y2;
            x1 = Integer.parseInt(st.nextToken());
            y1 = Integer.parseInt(st.nextToken());
            x2 = Integer.parseInt(st.nextToken());
            y2 = Integer.parseInt(st.nextToken());

            int result = pSum[x2][y2] - pSum[x1-1][y2] - pSum[x2][y1-1] + pSum[x1-1][y1-1];
            bw.write(result + "\n");
        }

        bw.flush();
        bw.close();
    }

}

시간복잡도

누적합을 계산할 때 $O(N^2)$이 걸리고, 구간합 계산은 M번 만큼 시행하므로 $O(M)$이 걸린다. 이 둘을 합치면 최종 시간복잡도는 $O(N^2+M)$이 된다.