[백준] 구간 합 구하기(2042)

문제

[백준] 구간 합 구하기(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(1 ≤ b ≤ N)번째 수를 c로 바꾸고 a가 2인 경우에는 b(1 ≤ b ≤ N)번째 수부터 c(b ≤ c ≤ N)번째 수까지의 합을 구하여 출력하면 된다.

입력으로 주어지는 모든 수는 -2^63보다 크거나 같고, 2^63-1보다 작거나 같은 정수이다.(2의 63승 입니다. velog 윗첨자 적용이 안되네요..)

출력

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

해설

초기화(init), 수정(update), 합(sum) 등의 메서드를 구현해야 하는 문제이기 때문에 세그먼트 트리 알고리즘을 공부하기에 가장 좋은 문제라고 생각한다.

풀이

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

public class Main {
	public static void main(String[] args) throws IOException {
		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 K = Integer.parseInt(st.nextToken());

		long[] nums = new long[N + 1];

		for (int i = 1; i <= N; i++) {
			nums[i] = Long.parseLong(br.readLine());
		}

		int h = (int) Math.ceil(Math.log(N) / Math.log(2));
		int segSize = (int) Math.pow(2, h + 1);

		long[] tree = new long[segSize];

		init(nums, tree, 1, 1, N);
		
		StringBuilder sb = new StringBuilder();
		
		for (int i = 0; i < M + K; i++) {
			st = new StringTokenizer(br.readLine(), " ");
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());
			long c = Long.parseLong(st.nextToken());

			if (a == 1) { // update
				long diff = c - nums[b];
				nums[b] = c;
				
				update(nums, tree, 1, 1, N, b, diff);
			} else if(a == 2) {
				sb.append(sum(tree, 1, 1, N, b, (int)c));
				sb.append("\n");
			}
		}
		
		System.out.println(sb.toString());
	}

	public static long sum(long[] tree, int node, int start, int end, int left, int right) {
		if (left > end || right < start)
			return 0;

		if (left <= start && end <= right) {
			return tree[node];
		}

		int mid = (start + end) / 2;
		return sum(tree, node * 2, start, mid, left, right) + sum(tree, node * 2 + 1, mid + 1, end, left, right);
	}

	public static void update(long[] nums, long[] tree, int node, int start, int end, int index, long diff) {
		if (index < start || index > end)
			return;

		tree[node] += diff;

		if (start != end) { // 리프노드가 아니라면 자식 노드 값도 바꿔줘야하니
			int mid = (start + end) / 2;
			update(nums, tree, node * 2, start, mid, index, diff);
			update(nums, tree, node * 2 + 1, mid + 1, end, index, diff);
		} 
	}

	public static long init(long[] nums, long[] tree, int node, int start, int end) {
		if (start == end)
			return tree[node] = nums[start];

		int mid = (start + end) / 2;

		return tree[node] = init(nums, tree, node * 2, start, mid) + init(nums, tree, node * 2 + 1, mid + 1, end);
	}
}