백준 11658: 구간 합 구하기 3

문제 요약

N×N개의 수가 N×N 크기의 표에 채워져 있다. 그런데 중간에 수의 변경이 빈번히 일어나고 그 중간에 어떤 부분의 합을 구하려 한다. 표의 i행 j열은 (i, j)로 나타낸다. (x1, y1)부터 (x2, y2)까지 합이란 x1 ≤ x ≤ x2, y1 ≤ y ≤ y2를 만족하는 모든 (x, y)에 있는 수의 합이다.

표에 채워져 있는 수와 변경하는 연산과 합을 구하는 연산이 주어졌을 때, 이를 처리하는 프로그램을 작성하시오.

문제 분류

• 자료 구조
• 세그먼트 트리
• 누적합
• 다차원 세그먼트 트리

문제 풀이

세그먼트 트리를 이용하여 풀 수도 있지만, 펜윅트리를 이용해서 풀어보았다. 펜윅트리는 세그먼트 트리보다 공간복잡도가 낮고, 구현도 어렵지 않다. 다만, 구간의 최댓값, 최솟값 등은 펜윅트리로 구할 수 없다.

펜윅트리는 누적합을 구간으로 나누어 갱신을 유연하게 만든 것 같았다. 펜윅트리의 배열 정의도 0~i까지의 누적합을 반환하기 때문이다. 즉, p[j] - p[i]를 통해 구간의 누적합을 구할 수 있다. 이를 2차원으로 확장하여 구현하였다.

풀이 코드

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

public class Main {
	static int n;
	static long f[][];
	
	static void add(int i, int j, long val) {
		val -= sum(i, j) - sum(i - 1, j) - sum(i, j - 1) + sum(i - 1, j - 1);
		for(int ii = i; ii <= n; ii += (ii & -ii)) {
			for(int jj = j; jj <= n; jj += (jj & -jj)) {
				f[ii][jj] += val;
			}
		}
	}
	
	static long sum(int i, int j) {
		long sum = 0;
		for(int ii = i; ii >= 1; ii -= (ii & -ii)) {
			for(int jj = j; jj >= 1; jj -= (jj & -jj))
				sum += f[ii][jj];
		}
		return sum;
	}
	public static void main(String[] args) throws Exception {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		StringBuilder sb = new StringBuilder();
		n = Integer.parseInt(st.nextToken());
		int m = Integer.parseInt(st.nextToken());
		f = new long[n + 1][n + 1];
		
		for(int i = 1; i <= n; i++) {
			st = new StringTokenizer(br.readLine());
			for(int j = 1; j <= n; j++) {
				add(i, j, Integer.parseInt(st.nextToken()));
			}
		}
		int command, x1, y1, x2, y2, val;
		
		while(m-- > 0) {
			st = new StringTokenizer(br.readLine());
			command = Integer.parseInt(st.nextToken());
			if(command == 1) {
				x1 = Integer.parseInt(st.nextToken()) - 1;
				y1 = Integer.parseInt(st.nextToken()) - 1;
				x2 = Integer.parseInt(st.nextToken());
				y2 = Integer.parseInt(st.nextToken());
				long area = sum(x2, y2) - sum(x2, y1) - sum(x1, y2) + sum(x1, y1);
				sb.append(area).append("\n");
			}
			else {
				x1 = Integer.parseInt(st.nextToken());
				y1 = Integer.parseInt(st.nextToken());
				val = Integer.parseInt(st.nextToken());
				add(x1, y1, val);
			}
		}
		System.out.print(sb);
	}
}