[JAVA] 백준 2042 세그먼트 트리

 

 

 

2042번: 구간 합 구하기

문제

어떤 N개의 수가 주어져 있다. 그런데 중간에 수의 변경이 빈번히 일어나고 그 중간에 어떤 부분의 합을 구하려 한다. 만약에 1,2,3,4,5 라는 수가 있고, 3번째 수를 6으로 바꾸고 2번째부터 5번째까지 합을 구하라고 한다면 17을 출력하면 되는 것이다. 그리고 그 상태에서 다섯 번째 수를 2로 바꾸고 3번째부터 5번째까지 합을 구하라고 한다면 12가 될 것이다.

풀이

수들이 주어졌을 때 빈번하게 수의 변경이 이루어지고, 중간 중간에 구간 합을 출력하시오.

세그먼트 트리를 활용하여 문제를 해결하였다.

📌 구간 합 세그먼트 트리

        int k = 0;
        while (true){
            if(Math.pow(2, k) >= N) break;
            k++;
        }

        int start_index = (int) Math.pow(2, k);
        tree = new long[start_index * 2];

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

        for(int i = start_index-1; i>0; i--){
            tree[i] = tree[i * 2] + tree[i * 2 +1];
        }

$2^k ≥ N$ 을 만족하는 k의 최솟값을 구해준다.

그리고 트리를 다음과 같이 구성한다.

• $2^{k+1}$크기의 배열을 선언하여 준다.
• $2^k$ ~ $2^k + N$까지의 인덱스는 입력으로 받는 수들을 채워준다.
• 1 ~ $2^k-1$까지의 인덱스는 자식 노드의 수를 합친 값을 채워 구간 합 세그먼트 트리를 만든다.
  • tree[i] = tree[i * 2] + tree[i * 2 + 1]

📌 세그먼트 트리 노드 값 갱신

    static void changeNumber(int index, long value){
        long diff = value - tree[index];

        while (index > 0){
            tree[index] += diff;
            index /= 2;
        }
    }

• 새로 변경되는 값 - 현재 값을 이용하여 변경 값을 구한다.
• 갱신될 노드부터 부모로 이동하면서 모든 구간합들을 갱신하여준다.

📌 세그먼트 트리 구간 합

    static long rageSum(int start_index, int end_index){
        long sum = 0;

        while (start_index <= end_index){
            if(start_index % 2 == 1) sum += tree[start_index];
            if(end_index % 2 == 0) sum += tree[end_index];

            start_index = (start_index+1) /2;
            end_index = (end_index -1)/2;

        }
        return sum;
    }

세그먼트 트리의 구간 합은 다음의 과정을 통하여 계산한다.

1. start_index % 2 == 1이라면 구간 합에 더해준다.
2. end_index % 2 == 0이라면 구간 합에 더해준다.
3. start_index = (start_index + 1) /2로 갱신해준다.
4. end_index = (end_index - 1) / 2로 갱신해준다.

📌 전체 코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

class Main {
    static long[] 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 k = 0;
        while (true){
            if(Math.pow(2, k) >= N) break;
            k++;
        }

        int start_index = (int) Math.pow(2, k);
        tree = new long[start_index * 2];

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

        for(int i = start_index-1; i>0; i--){
            tree[i] = tree[i * 2] + tree[i * 2 +1];
        }

        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()) + start_index-1;
            long c = Long.parseLong(st.nextToken());

            if(a == 1){
                changeNumber(b, c);
            }
            else if(a == 2){
                c += start_index-1;
                sb.append(rageSum(b, (int) c))
                        .append("\n");
            }
        }
        System.out.println(sb);
    }

    static void changeNumber(int index, long value){
        long diff = value - tree[index];

        while (index > 0){
            tree[index] += diff;
            index /= 2;
        }
    }

    static long rageSum(int start_index, int end_index){
        long sum = 0;

        while (start_index <= end_index){
            if(start_index % 2 == 1) sum += tree[start_index];
            if(end_index % 2 == 0) sum += tree[end_index];

            start_index = (start_index+1) /2;
            end_index = (end_index -1)/2;

        }
        return sum;
    }

}

결과