[BOJ] 백준 1275 커피숍2 (Java)

#1275 커피숍2

난이도 : 골드 1

유형 : 자료 구조 / 세그먼트 트리

1275번: 커피숍2

▸ 문제

모두 알다시피 동호는 커피숍의 마담이다. (마담이 무엇인지는 본인에게 물어보도록 하자.)

어느 날 커피숍의 손님 A씨가 동호에게 게임을 하자고 했다.

그 게임은 다음과 같은 규칙을 갖는다.

N개의 정수가 있으면, 동호는 다음과 같이 말한다. “3~7번째 수의 합은 무엇이죠?” 그러면 상대방은 “그 답은 000입니다. 그리고 8번째 수를 2로 고치도록 하죠” 그러면 동호는 “네 알겠습니다.”라고 한 뒤에 다시 상대방이 동호가 했던 것처럼 “8~9번째 수의 합은 무엇이죠?”라고 묻게된다. 이 것을 번갈아 가면서 반복하는 게임이다.

당신은 이 게임의 심판 역을 맡았다. 요컨대, 질문에 대한 답들을 미리 알아야 한다는 것이다.

당신의 머리가 출중하다면 10만개 가량 되는 정수와 10만턴 정도 되는 게임을 기억할 수 있을 것이다. 몇판 이 게임을 즐기던 동호는 많은 사람들이 이 게임을 하기를 바라게 되었고, 당신에게 심판 프로그램을 구현해달라고 요청했다.

▸ 입력

첫째 줄에 수의 개수 N과 턴의 개수 Q가 주어진다.(1 ≤ N, Q ≤ 100,000) 둘째 줄에는 처음 배열에 들어가 있는 정수 N개가 주어진다. 세 번째 줄에서 Q+2번째 줄까지는 x y a b의 형식으로 x~y까지의 합을 구하여라, a번째 수를 b로 바꾸어라 라는 뜻의 데이터가 주어진다.

입력되는 모든 수는 -231보다 크거나 같고, 231-1보다 작거나 같은 정수이다.

▸ 출력

한 턴마다 구한 합을 한 줄마다 한 개씩 출력한다.

 

문제 풀이  

구간 합을 구하는 문제로 세그먼트 트리를 사용하여 풀이할 수 있다.

• 선형 탐색으로 구하면 10만개의 쿼리를 10만개의 원소를 순차탐색해야 하므로 10만*10만의 O(Q*N)시간복잡도가 발생한다.
• 세그먼트 트리는 구간 합 쿼리를 트리의 특성을 이용하여 로그함수로 해결할 수 있다. 따라서 O(QlogN)으로 시간복잡도로 문제를 풀이 할 수 있다는 장점이 있다.

설계

1. 구간 합 트리 크기 구하기

노드의 시작점 1이므로 일반적으로 구하는 (트리 사이즈+1)을 해준다.

• 트리 사이즈(size) = 2^(트리 높이+1)-1
• 트리 높이(h) = log(n)/log(2)

static int getTreeSize(int n) {
	int h = (int)Math.ceil(Math.log(n)/Math.log(2));
	return (int) Math.pow(2, h+1); // +1
}

 

 

2. 구간 합 트리 초기화

구간 합 트리를 초기화 시켜준다. 원소 크기의 범위가 -2^31 ~ 2^31-1이므로 원소의 합은 int범위를 넘어가기 때문에 long으로 선언해줘야한다.

static long init(int s, int e, int node) {
	if(s == e) {
		return tree[node] = elements[s];
	}
	
	int mid = (s+e)/2;
	return tree[node] = init(s, mid, node*2)+init(mid+1, e, node*2+1);
}

 

[ 1 2 3 4 5 ] 구간 합 트리는 다음과 같이 이루어진다.

 

구간 합 트리 초기화

 

 

3. 구간 합 구하기 ( x~y)

x>y일 경우, x와 y의 위치를 바꿔준다. 구간합 트리를 구조를 이용하여 x~y가 위치하는 구간의 원소를 모두 더해준 다음 출력해주면 된다.

static long prefixSum(int s, int e, int node, int l, int r) {
	if(e < l || r < s) return 0;
	if(l <= s && e <= r) {
		return tree[node];
	}
	
	int mid = (s+e)/2;
	return prefixSum(s, mid, node*2, l, r)+ prefixSum(mid+1, e, node*2+1, l, r);
}

 

[ 1 2 1 4 5]의 idx 2~4의 구간 합을 구한다고하면 해당 로직은 다음 두 개의 노드를 조회하여 더해준 다음 10을 출력하게 된다.

구간 합 구하기

 

 

4. 구간 합 트리 업데이트 ( idx, dif)

dif는 변하게 되는 값으로 변경으로 주어진 값(value)에서 현재 원소의 크기를 뺀 값이다. 

• dif = value - elements[idx]

 

해당 원소가 들어간 모든 트리의 노드를 dif값을 더해 업데이트 시켜준다.

static void update(int s, int e, int node, int idx, long dif) {
	if(s <= idx && idx <= e) {
		tree[node] += dif;
	}else return;
	
	if(s == e) return;
	
	int mid = (s+e)/2;
	update(s, mid, node*2, idx, dif);
	update(mid+1, e, node*2+1, idx, dif);
}

 

[ 1 2 3 4 5] 구간 합 트리를 [ 1 2 1 4 5] 로 바꾼다고 하면 다음과 같이 업데이트 된다.

 

구간 합 트리 업데이트

 

풀이 코드 

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

public class Main {

	static long[] elements, tree;
	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 q = Integer.parseInt(st.nextToken());
		
		elements = new long[n];
		st = new StringTokenizer(br.readLine());
		for(int i=0; i<n; i++) {
			elements[i] = Integer.parseInt(st.nextToken());
		}
		int size = getTreeSize(n);
		tree = new long[size];
		init(0,n-1,1);

		StringBuilder sb = new StringBuilder();
		for(int i=0; i<q; i++) {
			st = new StringTokenizer(br.readLine());
			int x = Integer.parseInt(st.nextToken())-1;
			int y = Integer.parseInt(st.nextToken())-1;
			
			if(x>y) {
				int tmp =x;
				x = y;
				y = tmp; 
			}
			int idx = Integer.parseInt(st.nextToken())-1;
			int value = Integer.parseInt(st.nextToken());
			
			long pSum = prefixSum(0, n-1, 1, x, y);
			sb.append(pSum+"\n");
			
			long dif = value - elements[idx];
			update(0, n-1, 1, idx, dif);
			elements[idx] = value;
		}
		System.out.println(sb.toString());
		
	}
	
	static int getTreeSize(int n) {
		int h = (int)Math.ceil(Math.log(n)/Math.log(2));
		return (int) Math.pow(2, h+1);
	}
	
	static long init(int s, int e, int node) {
		if(s == e) {
			return tree[node] = elements[s];
		}
		int mid = (s+e)/2;
		return tree[node] = init(s, mid, node*2)+init(mid+1, e, node*2+1);
	}
	
	static void update(int s, int e, int node, int idx, long dif) {
		if(s <= idx && idx <= e) {
			tree[node] += dif;
		}else return;
		
		if(s == e) return;
		
		int mid = (s+e)/2;
		update(s, mid, node*2, idx, dif);
		update(mid+1, e, node*2+1, idx, dif);
	}
	
	static long prefixSum(int s, int e, int node, int l, int r) {
		if(e < l || r < s) return 0;
		if(l <= s && e <= r) {
			return tree[node];
		}
		
		int mid = (s+e)/2;
		return prefixSum(s, mid, node*2, l, r)+ prefixSum(mid+1, e, node*2+1, l, r);
	}
}