[백준 문제 풀기] 2042번: 구간 합 구하기 (Segment Tree, 자바)

문제 : https://www.acmicpc.net/problem/2042

 

문제
어떤 N개의 수가 주어져 있다. 그런데 중간에 수의 변경이 빈번히 일어나고 그 중간에 어떤 부분의 합을 구하려 한다. 만약에 1,2,3,4,5 라는 수가 있고, 3번째 수를 6으로 바꾸고 2번째부터 5번째까지 합을 구하라고 한다면 17을 출력하면 되는 것이다. 그리고 그 상태에서 다섯 번째 수를 2로 바꾸고 3번째부터 5번째까지 합을 구하라고 한다면 12가 될 것이다.

입력
첫째 줄에 수의 개수 N(1 ≤ N ≤ 1,000,000)과 M(1 ≤ M ≤ 10,000), K(1 ≤ K ≤ 10,000) 가 주어진다. M은 수의 변경이 일어나는 횟수이고, K는 구간의 합을 구하는 횟수이다. 그리고 둘째 줄부터 N+1번째 줄까지 N개의 수가 주어진다. 그리고 N+2번째 줄부터 N+M+K+1번째 줄까지 세 개의 정수 a, b, c가 주어지는데, a가 1인 경우 b번째 수를 c로 바꾸고 a가 2인 경우에는 b번째 수부터 c번째 수까지의 합을 구하여 출력하면 된다.

입력으로 주어지는 모든 수는 -2^63보다 크거나 같고, 2^63-1보다 작거나 같은 정수이다.

출력
첫째 줄부터 K줄에 걸쳐 구한 구간의 합을 출력한다. 단, 정답은 -2^63보다 크거나 같고, 2^63-1보다 작거나 같은 정수이다.

 

예제 입력, 출력

	예제 입력	예제 출력
1	5 2 2 1 2 3 4 5 1 3 6 2 2 5 1 5 2 2 3 5	17 12

 

사용 알고리즘

Segment Tree

 

코드 (자바)

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

// 구간 합 구하기 : https://www.acmicpc.net/problem/2042

public class Main {
	
	static int N = 0;	// 숫자의 개수
	static int M = 0;	// 수의 변경이 일어나는 횟수
	static int K = 0;	// 구간의 합을 구하는 횟수
	static long[] tree;	// 입력받은 트리
	
	static int init() {
		Arrays.fill(tree, 0);	// tree의 모든 값을 0 으로 초기화
		int ret = 1;
		while (ret<N) {
			ret*=2;
		}
		
		ret = ret-1;	// 1 base 구성을 위해서 1을 빼줌 
		
		return ret;
	}
	
	// 상위 노드로 올라가면서 구간합 업데이트
	static void update(long delta, int index) {
		while (index > 0) {
			tree[index] = tree[index] - delta;
			index/=2;
		}
	}
	
	// 구간합 출력
	static long sum(int from, int to) {
		long result = 0;
		
		while (from <= to) {
			if (from % 2 == 1) {result+=tree[from];}	// 시작점이 부모 노드 기준으로 오른쪽에 있는 자식일 때
			if (to % 2 == 0 ) {result+=tree[to];}	// 끝나는 점이 부모 노드 기준으로 왼쪽에 있는 자식일 때
			
			from = (from+1)/2;
			to = (to-1)/2;
		}
		
		return result;
		
	}
	
	public static void main(String[] args) throws Exception{
		// 입력 값 처리
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		
		N = Integer.parseInt(st.nextToken());
		M = Integer.parseInt(st.nextToken());
		K = Integer.parseInt(st.nextToken());
		
		tree = new long[1000000*3];		// segment tree 구성을 위해서 입력받을 수 있는 최대 숫자 개수의 3배를 배열로 선언
//		tree = new long[50*3];		// 완료 후 위에꺼로 바꿔야 함 (테스트용)
		int sIndex = init();		// tree 초기화 및 입력받은 숫자의 시작 인덱스 값을 구함
		
		// 노드에 입력받은 값 셋팅
		for (int i=1 ; i<=N ; i++) {
			tree[sIndex+i] = Integer.parseInt(br.readLine());
		}
		
//		System.out.println(Arrays.toString(tree));
		
		// 입력받은 값으로 상위노드에 합계를 채워 나감
		for (int i=sIndex ; i>0 ; i--) {
			tree[i] = tree[i*2] + tree[i*2+1];
		}
		
//		System.out.println(Arrays.toString(tree));
		
		// 입력받은 연산 처리
		for (int i=1 ; i<=M+K ; i++) {
			st = new StringTokenizer(br.readLine());
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());
			int c = Integer.parseInt(st.nextToken());
			
			if (a == 1) {
				// a가 1인 경우 b번째 수를 c로 바꾸고
				update (tree[b+sIndex]-c, b+sIndex);	// 증분값과 Segment Tree 에서의 인덱스를 전달
				
//				System.out.println(Arrays.toString(tree));

			} else if (a == 2) {
				// a가 2인 경우에는 b번째 수부터 c번째 수까지의 합을 구하여 출력
				System.out.println(sum(b+sIndex, c+sIndex));				
			}
		}
	}
}