세그먼트 트리

세그먼트 트리는 데이터들의 구간 합이나 구간에서 가장 큰수, 가장 작은수등을 빠르게 계산하기 위해 고안된 자료구조이다.

 

세그먼트 트리가 어떻게 구성되는지 예를 통해 알아보자

 

 

길이가 8인 배열 {1, 2, 3, 4, 5, 6, 7, 8} 이 있을 때 만약 누적합을 세그먼트 트리를 이용해 계산하면
우선 트리에 가장 아래 부분에 순서대로 수열의 값을 저장한다

 

트리의 부모노드에 자식노드의 합을 저장한다.

 

여기에 1 부터 인덱스를 부여하면

길이가 16인 배열에 { 0, 26, 10, 26, 3, 5, 7 ,11, 15, 1, 2, 3, 4, 5, 6, 7, 8} 순서대로 값이 저장되어 있다.

 

세그먼트 트리를 만드는 방법을 그림을 통해 유추해 보면

세그먼트 트리 배열의 길이는 기존의 배열의 길이보다 크거나 같은 2의 배수를 우선 구한 후에 곱하기 2를 하면 구할수 있는데

예를들어 기존배열의 길이가 3이면 가장 크면서 가까운 2의 배수는 4, 여기서 2를 곱한 8이 세그먼트 트리 배열의 크기이고

그림의 예처럼 배열의 길이가 8이면 크거나 같은 8에서 2를 곱한 16이 된다.

대충 기존배열 길이에 4를 곱하기도 한다

 

다시 그림을 보도록 하자

기존 배열은 인덱스 8부터 15번에 저장되어 있다. 따라서 기존배열은 세그먼트 트리에서 8부터 시작한다.

그리고 인덱스 7번에는 인덱스14와 인덱스 15의 값이 저장되어 있다.

6번 인덱스에는 인덱스 12와 인덱스 13의 값이 저장되어 있다.

이런식으로 부분의 합이 트리에 저장되어있는것을 확인할 수 있다.

 

이제 부분합을 구해 보도록 하자.

기존 배열 {1, 2, 3, 4, 5, 6, 7, 8} 에서 2번째 부터 7번째 (0부터 시작하는 인덱스의 경우 1번 인덱스부터 6번 인덱스) 까지 부분합을 구해보면

노란색 동그라미 쳐진 부분을 다 더하면 원하는 결과를 얻을 수가 있다.

세그먼트 트리에서 노란색 부분만 찾아내는것이 다소 어려울수 있는데 

우선 기억해야 할 것은

startIndex 를 2로 나누었을때 1이 남으면 해당값만 더하고 startIndex 에 1을 더하고

endIndex 를 2로 나누었을때 0이 남으면 해당값만 더하고 endIndex 에 1을 뺀다

이게 무슨말이냐??

그림에서 startIndex 는 9인것을 알수 있다 (기존인덱스 1에서 8을 더함)

9 는 2로 나눴을때 1이 남는다. 
그 말은 9의 부모값에 저장된 값은 원하는 범위의 값이 아니라는 말이다.

다시 그림을 보면 9번 인덱스의 부모 트리는 4번 인덱스이다

4번 인덱스는 8번과 9번 인덱스의 합인데
8번 인덱스는 해당 범위에 들어가 있지 않다.

따라서 9번 인덱스에 값만 더하고 인덱스를 한칸 옮겨줘서 다음 연산을 진행해야 한다.

이번에 endIndex 를 보면 endIndex 에서 2를 나누면 0이 남는다.

endIndex 는 반대로 2로 나누어 떨어지면 그 부모트리에 해당하는 값은 범위를 넘어서는 값을 포함하고 있어서

endIndex 에 있는 값만 더하고 endIndex 에 -1을 해줘야 다음 계산을 계속할 수 있다.

위의 계산을 마친 후에 부모트리로 이동하기 위해 startIndex 와 end 인덱스를 2로 나누어고 위의 연산을 반복하면 원하는 값을 얻을 수 있다.

// 구간합 계산
            long result = 0;
            while (start_index <= end_index) {
                if (start_index % 2 == 1) result += tree[start_index++];
                if (end_index % 2 == 0) result += tree[end_index--];
                start_index /= 2;
                end_index /= 2;
            }

 

# 문제

1275번 : 커피숍

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

위 문제로 연습을 해 보자

중간에 값을 변경해 주는 절차가 있다.

 

그림으로 나타내면 다음과 같다

문제의 예제를 따라가 보자

주어진 배열 {1, 2, 3, 4, 5}

2에서 3 범위의 합을 구하라고 했다.

세그먼트트리에서 기존배열이 저장되어있는 주소는 8번 부터이므로 배열 2, 3의 값은 그림에서 보면 9번인덱스부터 10번인덱스에 저장되어 있다.

따라서 startIndex = 9, endIndex = 10 이라는 것을 알 수 있다.

startIndex 를 2로 나눴을때 나누어 떨어지면 부모트리로 이동하면 되지만 2로 나누어 떨어지지 않는다면 해당값을 더한후에 startIndex 에 1을 더한후 부모트리로 이동시켜야 한다

반대로 endIndex 를 2로 나눴을때 나누어 1이 남으면 부모트리로 이동하고 그렇지 않으면 해당값을 더한 후에 1을 뺸후 부모트리로 이동한다

// 구간합 계산
            long result = 0;
            while (start_index <= end_index) {
                if (start_index % 2 == 1) result += tree[start_index++];
                if (end_index % 2 == 0) result += tree[end_index--];
                start_index /= 2;
                end_index /= 2;
            }

그림을 보면서 해당 코드를 머릿속으로 돌려보면 답이 result 에 저장되는것을 알수 있다.

 

그리고 문제에서 따라서 값을 변경하면 3번재 값을 1로 변경했다고 하면

이렇게 트리의 결과가 바뀌게 된다.

            while (point > 0) {         // 부모 노드 갱신
                point /= 2;
                tree[point] = tree[point * 2] + tree[point * 2 + 1];
            }

 

해당 값을 변경하고 인덱스 / 2를 한후에 값을 순서대로 갱신하면 트리의 값을 올바르게 최신화 할 수 있다.

 

전체 코드

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

public class Main {
    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 start = 1;
        while (start < N) start *= 2; // 리프 노드 개수 설정
        long[] tree = new long[start * 2]; // 세그먼트 트리 배열
        
        // 입력 배열 초기화
        st = new StringTokenizer(br.readLine());
        for (int i = 0; i < N; i++) tree[start + i] = Integer.parseInt(st.nextToken());
        
        // 내부 노드 초기화
        for (int i = start - 1; i > 0; i--) tree[i] = tree[i * 2] + tree[i * 2 + 1];
        
        // 쿼리 처리
        for (int i = 0; i < M; i++) {
            st = new StringTokenizer(br.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());
            
            // x와 y를 정렬하여 항상 x <= y가 되도록 설정
            int start_index = Math.min(x, y) + start - 1;
            int end_index = Math.max(x, y) + start - 1;

            // 구간합 계산
            long result = 0;
            while (start_index <= end_index) {
                if (start_index % 2 == 1) result += tree[start_index++];
                if (end_index % 2 == 0) result += tree[end_index--];
                start_index /= 2;
                end_index /= 2;
            }
            
            System.out.println(result);
            
            // 값 업데이트
            int point = a + start - 1; // 업데이트할 위치
            tree[point] = b;           // 새로운 값으로 변경
            while (point > 0) {         // 부모 노드 갱신
                point /= 2;
                tree[point] = tree[point * 2] + tree[point * 2 + 1];
            }
        }
    }
}