10 - 1. Prefix Sum 예제 (11660: 구간 합 구하기 5)

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

11660번: 구간 합 구하기 5

 

구간 합 구하기, 제목만 봐도 Prefix Sum을 사용해야 할 것 같은 문제입니다.

아닌 경우도 있습니다. ㅎㅎ

그 중에서도 2차원 Prefix Sum 문제 입니다.

 

왜 Prefix로 풀어야하는 지는 아래와 같습니다.

N은 최대 1024, M인 쿼리문은 최대 10만개까지 들어옵니다.

무식하고 돌리면 O (N * N * M) = 약 1000억, 따라서.. 일반 반복문으로 풀긴 어렵겠죠.

세그먼트나 인덱스 트리로 구현하기엔 각 행간 합은 어차피 반복문 N번 돌아야하니.. 좋은 방법은 아니겠습니다.

 

 

어쨌든, (x1, y1)부터 (x2, y2)까지 합을 구하는 프로그램을 구해라가 주된 목적입니다.

예시도 있는데, 대충 아래와 같은 이야기 입니다.

만약 쿼리로 2 1 3 3이 들어온다면, 아래의 범위 내 숫자들을 의미합니다.

그림 1. 예시 표

따라서 결과는 4 + 9 + 5 + 2 + 88 + 3 = 111을 출력하면 됩니다.

중첩 반복문은 불가능하다는 것을 알았고, prefix sum을 할 줄 아니까 이를 2차원으로 적용하면 되겠네요.

모르신다면, 여기로...

 

2차원 구조는 조금 독특한데, 살짝 파악해봅시다.

그림 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]는 어떻게 정의될 수 있을까요. 그림으로 살펴보겠습니다.

그림 3. sum[i][j]

sum[i][j]는 우선 arr[i][j]의 값은 가지고 있어야합니다.

그리고, 이전까지의 합들은 1차원과 같은 선형이 아니기 위의 그림 3처럼 모든 범위를 더하면 안됩니다.

선형으로 데이터를 더해 sum을 구하는 경우 위와 같은 범위 내 구간 합이 생성됩니다.

2차원에 맞는 부분적 합으로 구해야합니다.

 

 

부분들을 나누면 아래와 같습니다.

그림 4. 부분 분할

 

그림을 보면 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];
		}
	}
}

즉, 매번 이런식으로 배열에 값을 저장해주면, 빨간 부분을 제외하고 각 부분 합을 구할 수 있습니다.

이렇게 코드를 구성하는 것(빨간 블럭부분을 제외)이 문제의 키 포인트 입니다.

이해가 안되실 수 있는데, 직접 해보시는것을 추천드립니다. 그림으로 설명드리기엔 한계가 있거든요..

 

그래도 블럭이 생성되는 과정과 각 배열이 어느 범위를 의미하는지 몇 개만 보여드리면 아래와 같습니다.

그림 5. sum 배열 생성 과정

그림 6. 인덱스 별 의미

임의의 인덱스 (x1, y1) 일 때, 각각의 사각형은 아래와 같은 의미를 지닙니다.

 

초록색 사각형에 포함되는 원소의 합: sum[x1 + 3][y1]

노란색 사각형에 포함되는 원소의 합: sum[x1 + 4][y1 + 5]

주황색 사각형에 포함되는 원소의 합: sum[x1 + 2][y1 + 4]

 

 

이 정도면, 어떻게 되어있는지 어느 정도 와 닿으실 것 같습니다.

따라서 사각 범위의 누적합은 아래와 같이 구합니다.

그림 7. 타겟 범위

그림 8. 각 부분 블럭

target = A - B - C + D

D부분은 B, C 모두 포함하고 있기 때문에 한 번 더하기로 중복 연산을 없애줍니다.

 

그림 9. 구간 합 생성 과정

즉, 해당 블럭들을 그림 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만 조금 넘네요.

그림 10. 결과

 

sum 배열을 초기에 구성하는 방식이 1차원과는 사뭇 다릅니다.

DP 문제를 풀어보신들은 금방 이해하실 것 같은데요. 이해가 안된다면 귀찮더라도 직접 그려보시길 권장합니다.

 

 

그래도 이해가 안되신다면 문의해주세요. 좀 더 자세한 설명 드리도록 하겠습니다.