자료구조 공부/Segment Tree

백준 1275 ( 커피숍 2)

kdhoooon 2021. 2. 16. 14:57

문제


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

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

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

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

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

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

 

 

 

풀이


알고리즘 과 자료구조는 많은 이론을 모르면 풀 수 없다고 느낀 문제중 하나.

처음에는 반복문을 통해 풀었다. 그렇게 될경우 실행시간이 O(N * M) 으로 M이 커지면 커질수록 실행시간이 길어져서 시간 초과가 나온다.

<반복문 코드>

import java.io.*;
import java.math.*;
import java.security.*;
import java.text.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.function.*;
import java.util.regex.*;
import java.util.stream.*;
import static java.util.stream.Collectors.joining;
import static java.util.stream.Collectors.toList;

public class Main {	
	
	static StringBuilder sb = new StringBuilder();

	public static void main(String[] args) throws IOException{
		
		BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(System.in));
		
		StringTokenizer st = new StringTokenizer(bufferedReader.readLine());
		int n = Integer.parseInt(st.nextToken());
		int q = Integer.parseInt(st.nextToken());
		
		int[] list = new int[n + 1];
		st = new StringTokenizer(bufferedReader.readLine());
		for(int i = 1 ; i <= n ; i++) {
			
			list[i] = Integer.parseInt(st.nextToken());
		}
		
		int tmp;
		long total = 0;
		for( int i = 0 ; i < q ; i++) {
			st = new StringTokenizer(bufferedReader.readLine());
			
			int x = Integer.parseInt(st.nextToken());
			int y = Integer.parseInt(st.nextToken());
			
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());
			
			if(x > y) {
				tmp = x;
				x = y;
				y = tmp;
			}
			
			total = 0;
			for(int j = x ; j <= y ; j++) {
				total += list[j];
			}
			
			sb.append(total + "\n");
			
			list[a] = b;
		}
		
		
		System.out.println(sb);
	}
}

 

 

문제를 찾아보니 Segment Tree 를 이용해서 푸는 것이다 Segment Tree란 구간 합을 미리 구해 저장을 해놓은 트리를 말한다.

 

0 ~ 11 까지의 수를 가지고 있는 배열이 있다고 했을 때

이 와 같이 해당 숫자들을 저장 시켜 놓은 트리를 말한다.

 

트리를 1차배열로 구성하기 때문에

 

i 번째 트리의 자식은 i * 2 , i * 2 + 1 임을 이용하여 재귀적으로 문제를 풀면 된다.

 

처음 입력받은 배열을 통해 구간 합을 구해 트리를 구성하는 함수

public static long init(int start, int end, int node) {
	if(start == end) {
		return tree[node] = list[start];
	}
	
	int mid = (start + end) / 2;
	
	return tree[node] = init(start, mid, node * 2) + init(mid + 1, end, node * 2 + 1);
}

 

여기서 보면 mid 를 만들어 시작점과 끝점의 중간을 나눠가며 구간합을 미리 더해 놓는것을 알 수있다.

자식 노드의 인덱스는 해당 노드의 인덱스에서 2를 곱한값과, 2를 곱하고 1을 더한값으로 재귀적으로 접근하는 것을 알 수 있다.

 

구간합을 구하는 함수

public static long sum(int start, int end, int node, int x, int y) {
		
	if(end < x || y < start) {
		return 0;
	}
		
    if(x <= start && end <= y) {
		return tree[node];
	}
	
	int mid = (start + end) / 2;
	return sum(start, mid, node * 2, x , y) + sum(mid + 1, end, node*2 + 1, x , y);
}

 

a 노드의 값을 b 로 바꾸고 해당되는 노드의 구간합을 업데이트 해주는 함수

public static long update(int start, int end, int node, int a, int b) {
	
	if(a < start || a > end) {
		return tree[node];
	}
	
	if(start == end) {
		return tree[node] = b;
	}
	
	int mid = (start + end) / 2;
	return tree[node] = update(start, mid, node * 2, a, b) + update(mid + 1, end, node * 2 + 1, a, b);
}

 

 

<전체 코드>

import java.io.*;
import java.math.*;
import java.security.*;
import java.text.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.function.*;
import java.util.regex.*;
import java.util.stream.*;
import static java.util.stream.Collectors.joining;
import static java.util.stream.Collectors.toList;

public class Main {	
	
	static StringBuilder sb = new StringBuilder();
	static int[] list;
	static long[] tree;

	public static void main(String[] args) throws IOException{
		
		BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(System.in));
		
		StringTokenizer st = new StringTokenizer(bufferedReader.readLine());
		int n = Integer.parseInt(st.nextToken());
		int q = Integer.parseInt(st.nextToken());
		
		list = new int[n + 1];
		st = new StringTokenizer(bufferedReader.readLine());
		for(int i = 1 ; i <= n ; i++) {			
			list[i] = Integer.parseInt(st.nextToken());
		}
		
		tree = new long[n * 4];
		init(1, n, 1);
		
		int tmp;
		for( int i = 0 ; i < q ; i++) {
			st = new StringTokenizer(bufferedReader.readLine());
			
			int x = Integer.parseInt(st.nextToken());
			int y = Integer.parseInt(st.nextToken());
			
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());
			
			if(x > y) {
				tmp = x;
				x = y;
				y = tmp;
			}
			
			sb.append(sum(1, n, 1, x, y) + "\n");
			
			update(1, n, 1, a, b);
		}
		
		
		System.out.println(sb);
	}
	
	public static long init(int start, int end, int node) {
		if(start == end) {
			return tree[node] = list[start];
		}
		
		int mid = (start + end) / 2;
		
		return tree[node] = init(start, mid, node * 2) + init(mid + 1, end, node * 2 + 1);
	}
	
	public static long sum(int start, int end, int node, int x, int y) {
		
		if(end < x || y < start) {
			return 0;
		}
		
		if(x <= start && end <= y) {
			return tree[node];
		}
		
		int mid = (start + end) / 2;
		return sum(start, mid, node * 2, x , y) + sum(mid + 1, end, node*2 + 1, x , y);
	}
	
	public static long update(int start, int end, int node, int a, int b) {
		
		if(a < start || a > end) {
			return tree[node];
		}
		
		if(start == end) {
			return tree[node] = b;
		}
		
		int mid = (start + end) / 2;
		return tree[node] = update(start, mid, node * 2, a, b) + update(mid + 1, end, node * 2 + 1, a, b);
	}
}

'자료구조 공부 > Segment Tree' 카테고리의 다른 글

백준 5676 (음주 코딩)  (0) 2021.03.30
백준 11505 ( 구간 곱 구하기 )  (0) 2021.03.05
백준 2357 (최솟값과 최댓값)  (0) 2021.03.05
백준 1168 (요세푸스 2)  (0) 2021.02.21