어떤 N개의 수가 주어져 있다. 그런데 중간에 수의 변경이 빈번하게 일어나고, 그 중간에 어떤 부분의 합을 구하려 한다. 만약 1, 2, 3, 4, 5라는 수가 있고, 3번째 수를 6으로 바꾸고 2번째부터 5번째까지 합을 구하라고 한다면 17을 출력하면 되는 것이다. 그리고 이 상태에서 5번째 수를 2로 바꾸고, 3번째부터 5번째까지 합을 구하라고 한다면 12가 될 것이다.
(시간 제한 2초)
입력
1번째 줄에 N(1 ≤ N ≤ 1,000,000)과 M(1 ≤ M ≤ 10,000), K(1 ≤ K ≤ 10,000)가 주어진다. N은 수의 개수, M은 수의 변경이 일어나는 횟수, K는 구간의 합을 구하는 횟수다. 그리고 2번째 줄부터 N + 1번째 줄까지 N개의 수가 주어진다. 그리고 N + 2번째 줄부터 N + M + K + 1번째 줄까지 3개의 정수 a, b, c가 주어지는데, a가 1일때 b(1 ≤ b ≤ N) 번째 수를 c 로 바꾸고, a가 2일 경우에는 b(1 ≤ b ≤ N)번째 수에서 c(b ≤ c ≤ N)번째 수까지의 합을 구해 출력하면 된다. 입력으로 주어지는 모든 수는 -263보다 크거나 같고, 263 - 1 보다 작거나 같은 정수다.
출력
1번째 줄부터 K줄에 걸쳐 구한 구간의 합을 출력한다. 단, 정답은 -263보다 크거나 같고, 263 - 1보다 작거나 같은 정수다.
예제 입력
5 2 2 // 노드 개수, 변경 횟수, 구간 합 횟수
1
2
3
4
5
1 3 6
2 2 5
1 5 2
2 3 5
예제 출력
17
12
1단계 - 문제 분석하기
단순하게 구간 합을 구하는 문제라면 합 배열 자료구조를 이용하면 되지만, 이 문제는 중간의 수의 변경이 일어난다. 따라서 이 문제는 데이터의 변경에도 시간이 비교적 적게 걸리는 세그먼트 트리 자료구조를 이용해 해결할 수 있다.
2단계 - 손으로 풀어 보기
1 1차원 배열을 이용해 트리의 값을 초기화한다. 트리 배열 크기가 N = 5 이므로 ≥ N을 만족하는 k의 값은 3이고, 배열의 크기는 * 2 = 16이 된다. 시작 인덱스는 = 8이 된다.
구간 합 공식 A[N] = A[2N] + A[2N + 1] 을 사용하여 - 1 인덱스부터 1번 쪽으로 배열을 채워간다.
2 질의값 연산 함수와 데이터 업데이트 함수를 수행하여 질의와 관련된 결괏값을 출력한다.
3단계 - sudo코드 작성하기
tree(세그먼트 트리 배열)
N(수의 개수) M(변경이 일어나는 횟수) K(구간 합을 구하는 횟수)
Math.pow(2, k) ≥ N을 만족하는 k의 값 구하기
트리 배열의 크기 = Math.pow(2, k) * 2
시작 인덱스 = Math.pow(2, k)
for(i -> N만큼 반복)
{
tree[i+시작 인덱스]에 데이터 입력받기
}
for(i = 시작 인덱스 - 1; i > 0; i--)
{
tree[i] = tree[i * 2] + tree[(i * 2) + 1];
}
for(M + K 만큼 반복)
{
a(질의 유형) s(시작 인덱스) e(변경값 또는 종료 인덱스)
if(a == 1)
changeVal(tree 시작 인덱스, e)
else if(a == 2)
getSum(tree 시작 인덱스, tree 종료 인덱스)
}
changeVal(시작 인덱스, 변경값)
{
diff(현재 노드 값과 변경된 값의 차이)
while(시작 인덱스 > 0)
{
시작 인덱스 트리값에 diff값을 더함
시작 인덱스 = 시작 인덱스 / 2
}
}
getSum(시작 인덱스, 종료 인덱스)
{
구간 합 = 0
while(시작 인덱스 <= 종료 인덱스)
{
if(시작 인덱스 % 2 == 1)
{
구간 합 += 해당 노드 값
시작 인덱스++
}
if(종료 인덱스 % 2 == 0)
{
구간 합 += 해당 노드 값
종료 인덱스--
}
시작 인덱스 = 시작 인덱스 / 2
종료 인덱스 = 종료 인덱스 / 2
}
}
4단계 - 코드 구현하기
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;
public class Q71 {
public static int[] 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 M = Integer.parseInt(st.nextToken());
int K = Integer.parseInt(st.nextToken());
int treeHeight = 1;
int k = 0;
while(treeHeight < N){
treeHeight *= 2;
k ++;
}
int treeSize = (int)Math.pow(2, k) * 2;
int start_index = (int)Math.pow(2, k);
tree = new int[treeSize];
for(int i = 0; i < N; i++){
tree[i + start_index] = Integer.parseInt(br.readLine());
}
for(int i = start_index-1; i > 0; i--){
tree[i] = tree[i * 2] + tree[(i * 2) + 1];
}
for(int i = 0; i < M + K; i++){
st = new StringTokenizer(br.readLine());
int a = Integer.parseInt(st.nextToken());
int s = Integer.parseInt(st.nextToken());
int e = Integer.parseInt(st.nextToken());
if(a == 1){
changeVal(s + (int)Math.pow(2, k) - 1 ,e);
}else if(a == 2){
System.out.println(getSum(s + (int)Math.pow(2, k) - 1, e + (int)Math.pow(2, k) - 1));
}
}
br.close();
}
public static void changeVal(int start_index, int value){
int diff = value - tree[start_index];
while (start_index > 0){
tree[start_index] += diff;
start_index = start_index / 2;
}
}
public static int getSum(int start_index, int end_index){
int sum = 0;
while (start_index <= end_index){
if(start_index % 2 == 1){
sum += tree[start_index];
start_index++;
}
if (end_index % 2 == 0){
sum += tree[end_index];
end_index--;
}
start_index = start_index / 2;
end_index = end_index / 2;
}
return sum;
}
}
- Do it! 알고리즘 코딩테스트 자바 편
