[BOJ - JAVA] 2042 - 구간 합 구하기 (세그먼트 트리)

# 주소

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

2042번: 구간 합 구하기

# 문제

# 문제 해설 및 코드 리뷰

import java.io.*;
import java.util.*;
public class Main {
    static int n, m,k;
    static long[] arr;
    static int count;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st = new StringTokenizer(br.readLine());
        n = Integer.parseInt(st.nextToken());
        m = Integer.parseInt(st.nextToken());
        k = Integer.parseInt(st.nextToken());
        for(count = 1; count < n; count *= 2);
        arr = new long[count * 2 + 2];

        for(int i = 1; i <= n; i++){
            arr[count + i - 1] = Long.parseLong(br.readLine());
        }
        for(int i = count-1; i >= 1; i--){
            arr[i] = arr[i*2] + arr[i*2+1];
        }

        int size = m + k;
        for(int i = 1; i <= size; i++){
            int type, x;
            long y;
            st = new StringTokenizer(br.readLine());
            type = Integer.parseInt(st.nextToken());
            x = Integer.parseInt(st.nextToken());
            y = Long.parseLong(st.nextToken());
            if(type == 1){
                edit(x, y);
            }
            else{
                int endId;
                endId = (int)y;
                bw.write(String.valueOf(sum(x,endId))+"\n");
            }
        }
        bw.flush();
        bw.close();
        br.close();
    }
    public static void edit(int id, long value){
        int x = id + count - 1;
        arr[x] = value;
        x = x / 2;
        while(x>0){
            arr[x] = arr[x*2] + arr[x * 2 +1];
            x = x / 2;
        }
        return;
    }
    public static long sum(int start, int end){
        int a = start + count - 1;
        int b = end + count - 1;
        long ret = 0;
        while(a <= b){
            if((a & 1) == 1)
                ret += arr[a++];
            if((b & 1) == 0)
                ret += arr[b--];
            a = a / 2;
            b = b / 2;
        }
        return ret;
    }
}

https://subbak2.tistory.com/17

[BOJ 백준] 구간합 구하기(2042) Java

세그먼트 트리에 대해서 제대로 공부했습니다.

투 포인터, 구간합 문제로도 불리는 이 유형은 난이도가 꽤나 있는 편에 속합니다.

주로 DFS, BFS, 백트래킹 위주로 풀어온 저는 세그먼트 트리 풀이에 적잖이 당황을 했었는데요..

여러 블로그, 백준 풀이, 유튜브 영상등을 찾아보며 제 것으로 만들기 위해 시간을 할애했습니다.

따라서 이 블로그 포스팅이 생각보다 많이 길어질 수 있겠지만 저의 공부를 위해서, 그리고 혹시나 모를 방문자들을 위해서 조금이라도 도움이 되었으면 좋겠다 싶어 미숙하지만 봐주시면 좋겠습니다.

 

일단 세그먼트 트리는, 일반적으로 구간 합 문제를 구할 때 주로 이용되는 기법입니다.

만약 주어지는 N값이 적다면 그냥 sum에 배열 요소들을 더해주어서 구간 합을 구해도 되겠지만 문제를 보시면 입력값이 2^63까지 주어진다 하니 웬만한 배열로는 풀다가 시간초과나기 일수입니다.

https://www.acmicpc.net/blog/view/9

세그먼트 트리 (Segment Tree)

문제에서는 1번 연산을 수행하면 배열 요소의 값을 바꾸기 때문에 O(N)의 시간 복잡도가 걸리고 2번 연산에서는 O(1)이 걸리며 이게 M번 수행해야 한다면 총 시간복잡도는 O(NM)이 나옵니다.

 

만약 일일이 수가 바뀔때마다 2번 연산을 수행해야 한다면 계속 구간합을 구하고 출력하고 구하고 출력하고를 반복해야 하기 때문에 시간 복잡도는 O(N)이 걸려 매우 비효율적인 풀이가 됩니다.

그러나 세그먼트 트리를 이용하면 1번 연산을 O(logN), 2번 연산을 O(logN)만에 수행할 수 있습니다.

이 세그먼트 트리를 이해하기 위해선 먼저 리프 노드와 다른 노드에 대해 이해할 필요가 있습니다.

• 리프 노드 : 배열의 그 수 자체
• 다른 노드 : 왼쪽 자식과 오른쪽 자식의 합을 저장

이렇게 이해해 봅시다.

그러면 N = 10인 경우에 세그먼트 트리는 다음과 같습니다.

또한 어떤 노드의 번호가 x일 때, 왼쪽 자식의 번호는 2 * x , 오른쪽 자식의 번호는 2 * x + 1입니다.

따라서 

이런 세그먼트 트리 구조를 가지게 될 것입니다.

그럼 n이 주어져 있을 때 총 노드의 개수는 어떻게 구할 수 있을까요?

자 주어지는 n에 따라 세그먼트 트리의 층이 높아질수도, 낮아질수도 있습니다.

따라서 2의 배수에 따라 1층, 2층, 3층씩,,, 점점 더 높아질 수 있겠지요.

그러므로 주어진 n이 만약 10이라면 2^k가 10보다 크려면 k가 4가 될때지요?

이 4는 루프의 높이가 4일 때를 이야기 하므로 가질 수 있는 경우의 수는 2^4 * 2인 32가 되겠습니다.

실제로 저 세그먼트 트리에 맥시멈으로 노드를 모두 다 채웠을 때 15번 노드의 자식 노드는 31, 32번이므로 총 노드의 개수는 32가 정답입니다.

위의 예제에서 N이 5일 때를 이야기하고 있으므로 2^k가 5보다 큰 수 중 가장 적은 수는 3이므로 세그먼트 크기는 16이 될 것입니다.

여기까진 이해 가셨나요??

다음으로 넘어가봅시다.

그럼 이제 세그먼트 트리에 값을 주입시켜봅시다.

 

세그먼트 트리 만들기

       for(count = 1; count < n; count *= 2);
        arr = new long[count * 2 + 2];

        for(int i = 1; i <= n; i++){
            arr[count + i - 1] = Long.parseLong(br.readLine());
        }
        for(int i = count-1; i >= 1; i--){
            arr[i] = arr[i*2] + arr[i*2+1];
        }

세그먼트 트리를 만들기 위해 값을 입력받고 구간합을 구해봅시다.

첫 번째 for문에서 count가 될 수 있는 최소 값을 구하여 봅시다.

예제에서 n이 5이므로 count는 4까지 사용이 되고 8이 나옵니다.

결과적으로 arr의 크기는 count가 8이므로 18의 크기를 갖습니다.

(16이라고 해도 되지만 나중에 더 살펴봐야하기 때문에 일단 넘어가겠습니다.)

그리고 현재 count는 8인데, 이 8을 기준으로 위, 아래를 나누어서 사용할 것입니다.

8부터 시작해서 8, 9, 10, 11, 12번 노드는 입력받는 데이터를 저장할 노드입니다. 즉, 제일 하단에 위치할 노드가 될 것입니다.

또한 8번 보다 적은 인덱스의 노드는 리프노드의 부모노드가 됩니다.

따라서 이제부터 이 구간합의 노드들을 미리 설정해봅시다.

우리가 아직 사용하지 않은 노드는 1~7번 노드와 13~18번 노드가 있습니다.

for(int i = count-1; i >= 1; i--){
    arr[i] = arr[i*2] + arr[i*2+1];
}

 이 for문에서, i = 7부터 시작하여 감소하면서 index를 배정하고 있습니다.

arr[7] = arr[14] + arr[15]가 되겠네요.

arr[6] = arr[12] + arr[13]이 되겠고, 쭉 가서

arr[1] = arr[2] + arr[3]까지 구현이 되겠네요.

따라서 노드로 표현하면 

이 형태에서 13,14,15번은 0인 상태가 되겠지요.

자 그럼 구간합까지도 다 구했습니다.

이제, 1번 연산이 실행될 떄와 2번 연산이 실행됐을 때 함수의 수정과 수정된 구간합을 구하는 메서드를 구현해봅시다.

int size = m + k;
for(int i = 1; i <= size; i++){
    int type, x;
    long y;
    st = new StringTokenizer(br.readLine());
    type = Integer.parseInt(st.nextToken());
    x = Integer.parseInt(st.nextToken());
    y = Long.parseLong(st.nextToken());
    if(type == 1){
        edit(x, y);
    }
    else{
        int endId;
        endId = (int)y;
        bw.write(String.valueOf(sum(x,endId))+"\n");
    }
}

이 메서드를 실행할 size는 m + k입니다. 

while문으로 입력하셔도 좋고, for문으로 size까지 돌려도 좋습니다.

type과 x를 입력 받는데, type은 1번 연산인지 2번 연산인지 구분할 거고 x는 인덱스의 위치, y는 end구간 또는 변경할 VALUE가 되겠습니다.

그리고 type, x, y를 입력받은 뒤 만약 type이 1이면 edit함수를, 2번이면 sum함수를 실행합니다.

1번 연산 먼저 살펴봅시다.

1번 연산인 edit 메서드

edit연산을 실행하기 위해서, 바꿀 인덱스의 위치인 id와 값인 long타입의 value를 받아옵니다.

이 때 노드의 진짜 위치는 실제 저장한 id의 위치에 count를 더하고 1을 뺀 값입니다.

public static void edit(int id, long value){
    int x = id + count - 1;
    arr[x] = value;
    x = x / 2;
    while(x>0){
        arr[x] = arr[x*2] + arr[x * 2 +1];
        x = x / 2;
    }
    return;
}

그래서 x 는 id + count - 1이라고 하겠습니다.

이후 새로 만든 x위치에 value를 삽입합니다.

그리고 x와 관련된 모든 arr의 구간합을 변경할텐데요.

만약 n이 5일 때 2번째 위치를 바꾸고자 한다면 

위 사진에서처럼 9번, 4번, 2번, 1번의 값들을 모두 대폭 수정해야 할 것입니다.

따라서 제일 처음 arr[x] 의 위치는 value로 넣고 이후 4번, 2번, 1번의 위치의 값을 수정하기 위해 x를 2씩 나누어줍니다.

단, 이 x가 0보다 클 때 까지에 한해서 말이죠.

이 구문은 while로 표현하는게 좋겠습니다. 언제까지 나눌지 모르니까요. (알 수는 있지만 while이 깔끔하죠?)

새로 구해진 구간합은 당연히 새로 구하는 것이 맞겠습니다.

그래서 arr[4]는 arr[8] + arr[9]이므로 갱신된 arr[9]를 넣음으로써 arr[4], arr[2], arr[1]을 모두 구해주고 return합니다.

이렇게 edit메서드는 종료됩니다.

그럼 2번 연산을 살펴볼까요?

 

2번 연산인 sum메서드

sum메서드는 기본적으로 start와 end를 입력받되 이것들은 모두 int타입입니다.

왜냐하면 둘 다 value가 아닌, Index이기 때문이지요.

public static long sum(int start, int end){
    int a = start + count - 1;
    int b = end + count - 1;
    long ret = 0;
    while(a <= b){
        if((a & 1) == 1)
            ret += arr[a++];
        if((b & 1) == 0)
            ret += arr[b--];
        a = a / 2;
        b = b / 2;
    }
    return ret;

자 a와 b를 새로 정의할 것입니다.

여기서 a의 값은 위의 edit메서드에서 x위치를 수정했을 때처럼 start + count - 1을 해준 값입니다.

b의 위치 역시 end + start - 1을 해준 것이 리프노드의 위치입니다.

long 타입의 let을 0으로 초기화하고 이것을 구간합으로 설정합니다.

while문으로 들어가서, a <= b인 동안 (이분 탐색이랑 비슷하죠??) a와 b의 위치에 따라 조금 코드를 수정해줄건데요.

if문 앞에 &있는거 보이시죠?

이것은 a와 1을 비트연산자 곱을 통해 만약 a가 1과 비트 연산했을 때 참이다 -> 이 말은 a가 홀수 인덱스를 가지고 있단 뜻입니다.

만약 예제에서 2번째부터 5번째 인덱스까지 구간합을 구하라는 문제가 나온다면 9번 노드부터 12번 노드까지 구해야합니다.

따라서 a는 9, b는 12가 되겠습니다.

이때 a와 1을 비트연산 시키면 이는 00001001이고 끝의 연산자가 1이므로 이는 첫 번째 for문에서 참이 됩니다.

이 때 ret에서는 9 그자체의 노드를 더해야 하므로 arr[9]를 더하고 a를 증가시켜 a는 10으로 만듭니다. 왜냐구요? 구간합이니까 a를 증가시켜주는 것입니다.

다음 if문으로 넘어와서 b는 현재 12번인데 1과 비트 연산을 하게 되면 0이 나옵니다.

그러므로 ret에는 12번 노드도 더해주고 b는 감소시킵니다.

여기까지 좋아요. 근데 그 다음부터 중요합니다.

자식노드에 대해서 다 구했으니 부모노드로 이동을 해야하기 때문입니다.

따라서 a는 5, b 역시 5번 노드가 됩니다.

이 5번 노드는 1과 비트연산을 했을 때 1이 나오므로 ret에는 arr[5]를 더하고 더이상 a는 b보다 작지 않으므로 while문은 종료됩니다.

결국 최종적으로 ret에는 9번노드, 12번 노드, 5번 노드가 더해졌습니다.

위의 노드 그림으로 다시 봅시다.

왜 5번노드가 생뚱맞게 등장했을까요?

이 그림을 보시면 이해가 빠릅니다.

우리가 구하고 싶은 것은 9번 노드, 10번노드, 11번 노드, 12번 노드의 합입니다.

근데 9번 노드는 단독으로 더해야 하기 때문에 더하고 나서 10번노드로 이동하고 2로 나누어 5가 됩니다.

12번노드의 경우도 단독으로 더해야 하기 때문에 더하고 나서 11번노드로 이동 후 2로 나누어 5가 됩니다.

더 이상 a 는 b보다 작지 않기 때문에 5번노드만 더해주고 끝이날텐데요.

이 5번노드의 구성을 잘 보시면 10번과 11번을 가지고 있습니다.

따라서 9~12번 노드의 합은 9번, 5번, 12번 노드만 가지고도 구간합을 구할 수 있는 것입니다.

이해 가셨나요?

다른 예시를 들어보겠습니다.

만약 1번에서 3번까지의 합은 어떨까요?

그럼 8번부터 10번 노드의 합을 구해야 할 것입니다.

근데 8번 노드는 1과 비트연산했을 때 0이므로 pass, 10번 노드는 1과 비트연산했을 때 0이므로 ret에 10번 노드를 더한 뒤 9가 되고 2로 나누면 4가 됩니다.

a역시 8번 노드에서 2로 나누어서 4가 되었고 4번 노드는 어떠할지 보겠습니다.

4번 노드는 1과 비트연산하게되면 0입니다.

따라서 ret에 4번 노드가 더해질 것입니다.

이 4번노드는 8번노드와 9번노드의 합이기 때문에 최종적으로 4번 노드 + 10번 노드 = 8번 + 9번 + 10번인 것입니다.

이처럼 세그먼트 트리를 이용하여 구간합을 미리 구한 상태에서 더 빠르게 구간합을 구할 수도 있습니다.

 

처음보는 개념이라서 정말 어려웠는데 나동빈님의 유튜브 영상과 각종 블로그, 백준 세그먼트 트리 정리를 보면서 많이 공부했습니다.

이것도 쓰는데 한 4시간 걸렸네요 ㅋㅋㅋ ㅠㅠ

그만큼 많이 배웠습니다.

세그먼트 트리를 투 포인터라고도 부릅니다.

담에도 세그먼트 트리로 좋은 문제를 소개시켜드리겠습니다.

감사합니다!