백준 11660 구간 합 구하기 5 Java: 2차원 누적합으로 O(1) 쿼리

백준 11660번을 처음 풀 때는 각 행의 누적합만 만들었다. 제출은 통과했지만, 원문에도 “이거 시간 초과가 나야 정상인 것 같은데”라고 적어 두었다. 최악에는 쿼리마다 최대 1,024개 행을 다시 훑으므로 약 1억 번의 덧셈이 필요했다.

• 백준 11660번 구간 합 구하기 5

통과 여부가 실행 환경의 여유에 좌우되지 않게 하려면 2차원 누적합을 만들어야 한다. 전처리는 O(N²), 각 직사각형 쿼리는 O(1)이 된다.

2차원 누적합의 기준을 먼저 정하기

prefix[row][col]을 (1, 1)부터 (row, col)까지의 합으로 정의한다. 현재 칸의 값을 value라고 하면 다음 식으로 채울 수 있다.

prefix[row][col]
= value
 + prefix[row - 1][col]
 + prefix[row][col - 1]
 - prefix[row - 1][col - 1]

위쪽 직사각형과 왼쪽 직사각형을 더하면 겹치는 왼쪽 위 영역이 두 번 포함된다. 마지막 항을 한 번 빼는 이유다.

배열의 0행과 0열은 0으로 비워 둔다. 그러면 첫 행이나 첫 열에서도 별도 분기 없이 같은 식을 사용할 수 있다.

직사각형 합을 네 항으로 구하기

쿼리가 왼쪽 위 (x1, y1), 오른쪽 아래 (x2, y2)를 준다고 하자. 원하는 합은 다음과 같다.

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

전체에서 위쪽과 왼쪽을 빼면 왼쪽 위가 두 번 빠진다. 그래서 마지막에 한 번 더한다.

원문에서는 x를 행, y를 열로 읽는 부분을 헷갈려 첫 제출에 실패했다고 적었다. 변수명을 row, col로 두고 식에서도 같은 순서를 유지하면 이 실수를 줄일 수 있다.

Java 코드

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class Main {
    public static void main(String[] args) throws Exception {
        BufferedReader reader = new BufferedReader(
                new InputStreamReader(System.in)
        );
        StringTokenizer tokenizer = new StringTokenizer(reader.readLine());
        int size = Integer.parseInt(tokenizer.nextToken());
        int queryCount = Integer.parseInt(tokenizer.nextToken());

        long[][] prefix = new long[size + 1][size + 1];

        for (int row = 1; row <= size; row++) {
            tokenizer = new StringTokenizer(reader.readLine());
            for (int col = 1; col <= size; col++) {
                long value = Long.parseLong(tokenizer.nextToken());
                prefix[row][col] = value
                        + prefix[row - 1][col]
                        + prefix[row][col - 1]
                        - prefix[row - 1][col - 1];
            }
        }

        StringBuilder answer = new StringBuilder();
        for (int query = 0; query < queryCount; query++) {
            tokenizer = new StringTokenizer(reader.readLine());
            int x1 = Integer.parseInt(tokenizer.nextToken());
            int y1 = Integer.parseInt(tokenizer.nextToken());
            int x2 = Integer.parseInt(tokenizer.nextToken());
            int y2 = Integer.parseInt(tokenizer.nextToken());

            long sum = prefix[x2][y2]
                    - prefix[x1 - 1][y2]
                    - prefix[x2][y1 - 1]
                    + prefix[x1 - 1][y1 - 1];
            answer.append(sum).append('\n');
        }

        System.out.print(answer);
    }
}

문제의 수치 범위에서는 int도 가능하지만, 누적합 코드를 다른 문제에 재사용할 때 값 범위를 다시 계산하지 않아도 되도록 long으로 두었다.

처음 풀이와 달라진 점

방식	전처리	쿼리 1개	전체
행별 누적합	O(N²)	O(N)	O(N² + MN)
2차원 누적합	O(N²)	O(1)	O(N² + M)

첫 풀이가 통과한 사실과 풀이가 충분히 안전한지는 별개의 문제였다. 당시 적었던 “요행을 바라는 건 실력이 아니다”라는 문장은, 지금 다시 보면 제한을 먼저 계산하고 그 제한을 보장하는 알고리즘을 택하자는 뜻으로 남는다.

누적합을 나머지 빈도로 바꾸는 문제는 백준 10986 나머지 합, 문자별 누적합은 백준 16139 인간-컴퓨터 상호작용에서 이어서 볼 수 있다.

검증 범위

Java source를 수동 검토했다. 1×1 행렬, 전체 범위, 한 칸 범위, 가장자리 범위를 포함한 작은 무작위 행렬을 직접 합산한 결과와 대조한다. 현재 환경에는 실제 JDK가 없어 compile·judge 재제출은 live 반영 전에 별도 확인이 필요하다.