10 - 1. Prefix Sum 예제 (11660: 구간 합 구하기 5)
11660번: 구간 합 구하기 5
첫째 줄에 표의 크기 N과 합을 구해야 하는 횟수 M이 주어진다. (1 ≤ N ≤ 1024, 1 ≤ M ≤ 100,000) 둘째 줄부터 N개의 줄에는 표에 채워져 있는 수가 1행부터 차례대로 주어진다. 다음 M개의 줄에는 네
www.acmicpc.net
구간 합 구하기, 제목만 봐도 Prefix Sum을 사용해야 할 것 같은 문제입니다.
아닌 경우도 있습니다. ㅎㅎ
그 중에서도 2차원 Prefix Sum 문제 입니다.
왜 Prefix로 풀어야하는 지는 아래와 같습니다.
N은 최대 1024, M인 쿼리문은 최대 10만개까지 들어옵니다.
무식하고 돌리면 O (N * N * M) = 약 1000억, 따라서.. 일반 반복문으로 풀긴 어렵겠죠.
세그먼트나 인덱스 트리로 구현하기엔 각 행간 합은 어차피 반복문 N번 돌아야하니.. 좋은 방법은 아니겠습니다.
어쨌든, (x1, y1)부터 (x2, y2)까지 합을 구하는 프로그램을 구해라가 주된 목적입니다.
예시도 있는데, 대충 아래와 같은 이야기 입니다.
만약 쿼리로 2 1 3 3이 들어온다면, 아래의 범위 내 숫자들을 의미합니다.
따라서 결과는 4 + 9 + 5 + 2 + 88 + 3 = 111을 출력하면 됩니다.
중첩 반복문은 불가능하다는 것을 알았고, prefix sum을 할 줄 아니까 이를 2차원으로 적용하면 되겠네요.
모르신다면, 여기로...
2차원 구조는 조금 독특한데, 살짝 파악해봅시다.
우선, 전체 색칠한 부분의 sum은 sum[x2][y2]가 되어야 할 겁니다. (파랑 + 노랑)
그리고 파란색 부분이 실제로 구할 부분이고, 나머지 색들은 파란색 부분에서 빼줌으로써 제외해야할 부분이죠.
이렇게 색칠하고 나니, 전체 합에서 부분합을 빼주면 되겠다~ 정도의 생각이 듭니다.
우선 어떻게 뺄지 고민은 차후에 하기로하고 sum 배열부터 정의해보겠습니다.
1차원 Prefix Sum에서 우리는 sum이라는 배열을 구했습니다.
sum [i] = arr[i] + sum[i - 1] 이렇게 정의된 배열이죠.
즉 0번 ~ i번 arr 배열의 값을 모두 더한 결과가 sum[i]로 저장됩니다.
그렇다면, 2차원에서 우리가 궁금한 것은 sum[i][j]가 될텐데요.
sum[i][j]는 어떻게 정의될 수 있을까요. 그림으로 살펴보겠습니다.
sum[i][j]는 우선 arr[i][j]의 값은 가지고 있어야합니다.
그리고, 이전까지의 합들은 1차원과 같은 선형이 아니기 위의 그림 3처럼 모든 범위를 더하면 안됩니다.
선형으로 데이터를 더해 sum을 구하는 경우 위와 같은 범위 내 구간 합이 생성됩니다.
2차원에 맞는 부분적 합으로 구해야합니다.
부분들을 나누면 아래와 같습니다.
그림을 보면 sum[i][j]를 정의할 수 있습니다.
단, 말씀드렸지만 위의 그림에서 sum[i - 1][j + 1], sum[i - 1][j + 2] 이하 빨간 부분은 포함 되지 않도록 정의합니다.
해당 부분이 포함될 경우 2차원이 아닌 1차원 누적합과 다를바 없게 되거든요.
따라서, 이에 맞는 식은 아래와 같습니다.
sum[i][j] = sum[i][j - 1] + sum[i - 1][j] - sum[i - 1][j - 1] + arr[i][j]
코드로 보면 아래와 같습니다.
private static void initMap(int n, int[][] arr) {
sum = new int[n + 1][n + 1];
for(int row = 1; row < n + 1; row++) {
for(int col = 1; col < n + 1; col++) {
sum[row][col] = arr[row][col] + sum[row - 1][col] + sum[row][col - 1]
- sum[row - 1][col - 1];
}
}
}즉, 매번 이런식으로 배열에 값을 저장해주면, 빨간 부분을 제외하고 각 부분 합을 구할 수 있습니다.
이렇게 코드를 구성하는 것(빨간 블럭부분을 제외)이 문제의 키 포인트 입니다.
이해가 안되실 수 있는데, 직접 해보시는것을 추천드립니다. 그림으로 설명드리기엔 한계가 있거든요..
그래도 블럭이 생성되는 과정과 각 배열이 어느 범위를 의미하는지 몇 개만 보여드리면 아래와 같습니다.
임의의 인덱스 (x1, y1) 일 때, 각각의 사각형은 아래와 같은 의미를 지닙니다.
초록색 사각형에 포함되는 원소의 합: sum[x1 + 3][y1]
노란색 사각형에 포함되는 원소의 합: sum[x1 + 4][y1 + 5]
주황색 사각형에 포함되는 원소의 합: sum[x1 + 2][y1 + 4]
이 정도면, 어떻게 되어있는지 어느 정도 와 닿으실 것 같습니다.
따라서 사각 범위의 누적합은 아래와 같이 구합니다.
target = A - B - C + D
D부분은 B, C 모두 포함하고 있기 때문에 한 번 더하기로 중복 연산을 없애줍니다.
즉, 해당 블럭들을 그림 9의 설명을 참고해 배열 식으로 변경하면 아래와 같은 식이 나옵니다.
target = sum[x2][y2](A) - sum[x1 - 1][y2](B) - sum[x2][y1 - 1](C) + sum[x1 - 1][y1 - 1](D)
이러한 방식으로 사각 범위의 누적합을 O(1)로 바꿔준다면 해당 문제를 제한 시간 내에 해결할 수 있습니다.
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.StringTokenizer;
public class Main {
private static final String NEW_LINE = "\n";
private static int[][] sum;
public static void main(String[] args) throws Exception{
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[][] section = new int[N + 1][N + 1];
for(int i = 1; i < N + 1; i++) {
st = new StringTokenizer(br.readLine());
for(int j = 1; j < N + 1; j++) {
section[i][j] = Integer.parseInt(st.nextToken());
}
}
initMap(N, section);
StringBuilder sb = new StringBuilder();
while(M-- > 0) {
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());
sb.append(sum[x2][y2] - sum[x1 - 1][y2] - sum[x2][y1 - 1] + sum[x1 - 1][y1 - 1]).append(NEW_LINE);
}
System.out.print(sb.toString());
}
private static void initMap(int n, int[][] arr) {
sum = new int[n + 1][n + 1];
for(int row = 1; row < n + 1; row++) {
for(int col = 1; col < n + 1; col++) {
sum[row][col] = arr[row][col] + sum[row - 1][col] + sum[row][col - 1]
- sum[row - 1][col - 1];
}
}
}
}
해당 문제의 시간 복잡도는 O (N^2 + M) 입니다.
많이 잡아야 100만 조금 넘네요.
sum 배열을 초기에 구성하는 방식이 1차원과는 사뭇 다릅니다.
DP 문제를 풀어보신들은 금방 이해하실 것 같은데요. 이해가 안된다면 귀찮더라도 직접 그려보시길 권장합니다.
그래도 이해가 안되신다면 문의해주세요. 좀 더 자세한 설명 드리도록 하겠습니다.