세그먼트 트리 합구하기

# 문제

2268번 : 수들의 합 7

https://www.acmicpc.net/problem/2268

 

 

세그먼트 트리를 이용해야 제한시간에 맞게 풀수 있는 문제이다.

 

내용자체는 복잡하지 않아서 세그먼트 트리를 알면 충분히 풀수 있는 문제이지만 자바에서는 index 가 0에서 부터 시작되는데 반해 문제에서는 index 가 1부터 시작하기 때문에 혼란이 올수 있다.

또 세그먼트 트리에 저장되는 값이 int 가 넘을 수 있기 때문에 반드시 long 자료형으로 배열을 선언해 주어야 올바른 답을 구할 수 있다.

그리고 입력되는 인자가 반드시 처음값보다 나중값이 크다는 보장이 없어서 이 부분도 처리해 주어야 한다.

 

우선 만약에 문제에서 인덱스가 1번부터 시작하는 경우 배열의 크기를 세그먼트 트리의 형태에 맞게 선언해주고 나서
시작점을 1빼주면 문제에서 입력값이 모두 -1빼지는 효과가 나서 자바와 문제사이의 index 시작점 문제를 해결 할 수 있다.

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

public class Main {
	static long[] tree;
	static int start;
	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());
		start = 1;
		while (start < N) start *= 2;
		tree = new long[start * 2];
		start--;
		Arrays.fill(tree, 0);
		int M = Integer.parseInt(st.nextToken());
		for (int i = 0; i < M; i++) {
			st = new StringTokenizer(br.readLine());
			int query = Integer.parseInt(st.nextToken());
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());
			if (query == 0) sum(a,b);
			else modify(a, b);
		}
	}
	static void sum(int a, int b) {
		int s = Math.min(a, b) + start;
		int e = Math.max(a, b) + start;
		long result = 0;
		while (s <= e) {
			if (s % 2 == 1) result += tree[s++];
			if (e % 2 == 0) result += tree[e--];
			s /= 2;
			e /= 2;
		}
		System.out.println(result);
	}
	static void modify(int a, int b) {
		int idx = a + start;
		tree[idx] = b;
		while (idx > 1) {
			idx /= 2;
			tree[idx] = tree[idx * 2] + tree[idx * 2 + 1]; 
		}
	}
}

 

함수를 적당히 분리해서 코딩을 하니 좀더 깔끔한 느낌이 든다.